support multi loras

app.py

@@ -18,6 +18,9 @@ import boto3
 from io import BytesIO
 from datetime import datetime
 
+from diffusers import UNet2DConditionModel
+
+
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
@@ -27,7 +30,20 @@ login(token=HF_TOKEN)
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
+
+# unet = UNet2DConditionModel.from_pretrained(
+#     base_model,
+#     torch_dtype=torch.float16,
+#     use_safetensors=True,
+#     variant="fp16",
+#     subfolder="unet",
+# ).to("cuda")
+
+
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
+
+
+
 MAX_SEED = 2**32-1
 
 class calculateDuration:
@@ -70,7 +86,7 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
 
 
 @spaces.GPU
-def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):
+def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
@@ -82,7 +98,7 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
             width=width,
             height=height,
             generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
+            cross_attention_kwargs={"scale": 1.0},
             max_sequence_length=256
         ).images[0]
         
@@ -90,21 +106,34 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
     return generate_image
 
 
-def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, lora_strings, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     
-    # with calculateDuration("Unloading LoRA"):
-    #     pipe.unload_lora_weights()
-        
-    # Load LoRA weights
-    if lora_repo and lora_name:
-        with calculateDuration(f"Loading LoRA weights for {lora_repo} {lora_name}"):
-            pipe.load_lora_weights(lora_repo, weight_name=lora_name)
 
+    # Load LoRA weights
+    if lora_strings:
+        with calculateDuration(f"Loading LoRA weights for {lora_strings}"):
+            pipe.unload_lora_weights()
+            lora_array = lora_strings.split(',')
+            adapter_names = []
+            for lora_string in lora_array:
+                parts = lora_string.split(':')
+                if len(parts) == 3:
+                    lora_repo, weights, adapter_name = parts
+                    # 调用 pipe.load_lora_weights() 方法加载权重
+                    pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                    adapter_names.append(adapter_name)
+                else:
+                    print(f"Invalid format for lora_string: {lora_string}")
+            
+            adapter_weights = [lora_scale] * len(adapter_names)
+            # 调用 pipeline.set_adapters 方法设置 adapter 和对应权重
+            pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+        
     # Set random seed for reproducibility
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
     
     if upload_to_r2:
         with calculateDuration("upload r2"):
@@ -131,8 +160,7 @@ with gr.Blocks(css=css) as demo:
         
         with gr.Column():
             prompt = gr.Text(label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False)
-            lora_repo = gr.Text( label="Repo", max_lines=1, placeholder="Enter a lora repo", visible=True)        
-            lora_name = gr.Text( label="Weights", max_lines=1, placeholder="Enter a lora weights",visible=True)
+            lora_strings = gr.Text( label="lora_strings", max_lines=1, placeholder="Enter a lora strings", visible=True)        
             run_button = gr.Button("Run", scale=0)
 
             with gr.Accordion("Advanced Settings", open=False):
@@ -164,7 +192,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = run_lora,
-        inputs = [prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket],
+        inputs = [prompt, cfg_scale, steps, lora_strings, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket],
         outputs=[result, seed, json_text]
     )