Update app.py

app.py

@@ -163,11 +163,17 @@ app.prepare(ctx_id=0, det_size=(640, 640))
 
 # download checkpoints
 print("Downloading checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/controlnet/config.json", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/controlnet/diffusion_pytorch_model.safetensors", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/ip-adapter.bin", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="image_encoder/pytorch_model.bin", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="image_encoder/config.json", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="checkpoint_105000/controlnet/config.json", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="checkpoint_105000/controlnet/diffusion_pytorch_model.safetensors", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="checkpoint_105000/ip-adapter.bin", local_dir="./checkpoints")
+
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="image_encoder/pytorch_model.bin", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="image_encoder/config.json", local_dir="./checkpoints")
+
+# Download Lora weights
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="LoRAs/3D_avatar/pytorch_lora_weights.safetensors", local_dir=".")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="LoRAs/coloringbook/pytorch_lora_weights.safetensors", local_dir=".")
+hf_hub_download(repo_id="briaai/BRIA-2.3-ID_Preservation", filename="LoRAs/One_line_portraits_Light/pytorch_lora_weights.safetensors", local_dir=".")
 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -176,6 +182,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 face_adapter = f"./checkpoints/checkpoint_105000/ip-adapter.bin"
 controlnet_path = f"./checkpoints/checkpoint_105000/controlnet"
 base_model_path = f'briaai/BRIA-2.3'
+lora_base_path = f"./LoRAs"
+
 resolution = 1024
 
 # Load ControlNet models
@@ -206,13 +214,19 @@ pipe.load_ip_adapter_instantid(face_adapter)
 
 clip_embeds=None
 
+Loras_dict = {
+        "":"",
+        "One_line_portraits_Light": "An illustration of ",
+        "3D_avatar": "An illustration of ",
+        "coloringbook": "An illustration of "
+    }
 
-
+lora_names = Loras_dict.keys()
 
 @spaces.GPU
-# def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name, lora_scale, progress=gr.Progress(track_tqdm=True)):
-def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, progress=gr.Progress(track_tqdm=True)):
-    # global CURRENT_LORA_NAME  # Use the global variable to track LoRA
+def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name, lora_scale, progress=gr.Progress(track_tqdm=True)):
+# def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, progress=gr.Progress(track_tqdm=True)):
+    global CURRENT_LORA_NAME  # Use the global variable to track LoRA
 
     if image_path is None:
         raise gr.Error(f"Cannot find any input face image! Please upload a face image.")
@@ -239,9 +253,6 @@ def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_imag
         files = [
             ('file', ('image_name.jpeg', image_file, 'image/jpeg'))  # Specify file name, file-like object, and MIME type
         ]
-        # headers = {
-        # 'api_token': 'a10d6386dd6a11ebba800242ac130004'
-        # }
         headers = {
         'api_token': os.getenv('BRIA_RMBG_TOKEN')  # Securely retrieve the token
         }
@@ -269,7 +280,32 @@ def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_imag
                     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    full_prompt = prompt
+    # full_prompt = prompt
+    if lora_name != CURRENT_LORA_NAME:  # Check if LoRA needs to be changed
+        if CURRENT_LORA_NAME is not None:  # If a LoRA is already loaded, unload it
+            pipe.disable_lora()
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            print(f"Unloaded LoRA: {CURRENT_LORA_NAME}")
+
+        if lora_name != "":  # Load the new LoRA if specified
+            # pipe.enable_model_cpu_offload()
+            lora_path = os.path.join(lora_base_path, lora_name, "pytorch_lora_weights.safetensors")
+            pipe.load_lora_weights(lora_path)
+            pipe.fuse_lora(lora_scale)  
+            pipe.enable_lora()
+
+            # lora_prefix = Loras_dict[lora_name]
+
+            print(f"Loaded new LoRA: {lora_name}")
+
+        # Update the current LoRA name
+        CURRENT_LORA_NAME = lora_name
+
+    if lora_name != "":
+        full_prompt = f"{Loras_dict[lora_name]} + " " + {prompt}"
+    else:
+        full_prompt = prompt
 
     print("Start inference...")    
     images = pipe(
@@ -341,19 +377,21 @@ with gr.Blocks(css=css) as demo:
             placeholder="Enter your prompt here", 
             info="Describe what you want to generate or modify in the image."
             )
-            
+
+            lora_name = gr.Dropdown(choices=lora_names, label="LoRA", value="",  info="Select a LoRA name from the list, not selecting any will disable LoRA.")
+
             submit = gr.Button("Submit", variant="primary")
  
             with gr.Accordion(open=False, label="Advanced Options"):
                 num_steps = gr.Slider( 
-                    label="Number of sample steps",
+                    label="Number of diffusion steps",
                     minimum=1,
                     maximum=100,
                     step=1,
                     value=30,
                 )
                 guidance_scale = gr.Slider(
-                    label="Guidance scale",
+                    label="cfg scale",
                     minimum=0.1,
                     maximum=10.0,
                     step=0.1,
@@ -367,27 +405,33 @@ with gr.Blocks(css=css) as demo:
                     value=1,
                 )
                 ip_adapter_scale = gr.Slider(
-                    label="ip adapter scale",
+                    label="ID Adapter scale",
                     minimum=0.0,
                     maximum=1.0,
                     step=0.01,
                     value=0.8,
                 )
                 kps_scale = gr.Slider(
-                    label="kps control scale",
+                    label="lnmks ControlNet scale",
                     minimum=0.0,
                     maximum=1.0,
                     step=0.01,
                     value=0.6,
                 )
                 canny_scale = gr.Slider(
-                    label="canny control scale",
+                    label="canny ControlNet scale",
                     minimum=0.0,
                     maximum=1.0,
                     step=0.01,
                     value=0.4,
                 )
-
+                lora_scale = gr.Slider(
+                    label="LoRA scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.7,
+                )
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -409,7 +453,8 @@ with gr.Blocks(css=css) as demo:
             api_name=False,
         ).then(
             fn=generate_image,
-            inputs=[img_file, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale],
+            # inputs=[img_file, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale],
+            inputs=[img_file, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name, lora_scale],
             # outputs=[gallery]
             outputs=gallery
         )