changes

app.py

@@ -9,7 +9,6 @@ import logging
 from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForImage2Image
 from huggingface_hub import login
 from diffusers.utils import load_image
-#from lora_loading_patch import load_lora_into_transformer
 import time
 from datetime import datetime
 from io import BytesIO
@@ -27,28 +26,28 @@ from huggingface_hub import hf_hub_download
 from diffusers.quantizers import PipelineQuantizationConfig
 from diffusers import (FluxPriorReduxPipeline, FluxInpaintPipeline, FluxFillPipeline, FluxKontextPipeline, FluxPipeline)
 
+
 # Login Hugging Face Hub
 HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 import diffusers
 
+
 # init
 dtype = torch.bfloat16
 device = "cuda:0"
 
+
 base_model = "black-forest-labs/FLUX.1-Krea-dev"
 
-# pipeline_quant_config = PipelineQuantizationConfig(
-#     quant_backend="bitsandbytes_4bit",
-#     quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
-#     components_to_quantize=["transformer", "text_encoder_2"],
-# )
 
 txt2img_pipe = FluxKontextPipeline.from_pretrained(base_model, torch_dtype=dtype)
 txt2img_pipe = txt2img_pipe.to(device)
 
+
 MAX_SEED = 2**32 - 1
 
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -69,13 +68,14 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
         
+
 def safe_trim_for_clip(text: str, max_words: int = 77) -> str:
-    # 简单按词裁，不破坏主 prompt。你也可以做更智能的关键词抽取。
     tokens = re.split(r"\s+", text.strip())
     if len(tokens) <= max_words:
         return text
     return " ".join(tokens[:max_words])
 
+
 def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     with calculateDuration("Upload images"):
         connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
@@ -94,17 +94,16 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
         buffer.seek(0)
         s3.upload_fileobj(buffer, bucket_name, image_file)
         print("upload finish", image_file)
-        # start to generate thumbnail
+        
+        # Generate thumbnail
         thumbnail = image.copy()
         thumbnail_width = 256
         aspect_ratio = image.height / image.width
         thumbnail_height = int(thumbnail_width * aspect_ratio)
         thumbnail = thumbnail.resize((thumbnail_width, thumbnail_height), Image.LANCZOS)
         
-        # Generate the thumbnail image filename
         thumbnail_file = image_file.replace(".png", "_thumbnail.png")
         
-        # Save thumbnail to buffer and upload
         thumbnail_buffer = BytesIO()
         thumbnail.save(thumbnail_buffer, "PNG")
         thumbnail_buffer.seek(0)
@@ -112,9 +111,11 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
         print("upload thumbnail finish", thumbnail_file)
     return image_file
 
+
 def generate_random_4_digit_string():
     return ''.join(random.choices(string.digits, k=4))
 
+
 @spaces.GPU(duration=120)
 def run_lora(
     prompt, 
@@ -152,7 +153,6 @@ def run_lora(
         try:
             pipe.unload_lora_weights()
         except Exception as _:
-            # 某些版本上未加载时调用可能抛异常，忽略
             pass
 
     adapter_names = []
@@ -170,7 +170,6 @@ def run_lora(
                 for lora_info in lora_configs:
                     repo = lora_info.get("repo")
                     weights = lora_info.get("weights")
-                    # 优先使用用户提供的 adapter_name；没有则随机
                     adapter_name = lora_info.get("adapter_name") or f"adp_{generate_random_4_digit_string()}"
                     weight = float(lora_info.get("adapter_weight", 1.0))
                     if not (repo and weights):
@@ -178,7 +177,6 @@ def run_lora(
                         continue
                     try:
                         weight_path = hf_hub_download(repo_id=repo, filename=weights)
-                        # 关键修复：prefix=None，避免仅在 text_encoder 查找
                         pipe.load_lora_weights(weight_path, adapter_name=adapter_name, prefix=None)
                         adapter_names.append(adapter_name)
                         adapter_weights.append(weight)
@@ -187,8 +185,6 @@ def run_lora(
 
                 if adapter_names:
                     pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
-                    # 可选：融合后推理更快，但无法动态调整权重
-                    # pipe.fuse_lora(adapter_names=adapter_names)
                     try:
                         active = pipe.get_active_adapters() if hasattr(pipe, "get_active_adapters") else []
                         print("Active adapters:", active)
@@ -202,15 +198,15 @@ def run_lora(
                             lora_layer_count += 1
                     print(f"[DEBUG] transformer LoRA layers: {lora_layer_count}")
                     
-                    # 若层数为 0，给出直观警告
                     if lora_layer_count == 0:
                         gr.Warning("LoRA seems not injected (0 layers on transformer). Check whether the LoRA is trained for FLUX and `prefix=None` is set.")
 
-
     pipe.enable_vae_slicing()
     clip_side_prompt = safe_trim_for_clip(prompt, max_words=77) 
     init_image = None
     error_message = ""
+    image = None
+    
     try:
         gr.Info("Start to generate images ...")
         joint_attention_kwargs = {"scale": 1}
@@ -246,25 +242,25 @@ def run_lora(
     progress(100, "Completed!")
     
     # CHANGED: Return both image AND json
-    return image, json.dumps(result)  # <--- THIS IS THE KEY CHANGE
-
-# Gradio interface
-
+    return image, json.dumps(result)
 
 
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("flux-dev-multi-lora")
+    gr.Markdown("# flux-dev-multi-lora")
     with gr.Row():
         
         with gr.Column():
-
             prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=10)
-            lora_strings_json = gr.Text(label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1", "adapter_weight": 1}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2", "adapter_weight": 1}]', lines=5)
+            lora_strings_json = gr.Text(
+                label="LoRA Configs (JSON List String)", 
+                placeholder='[{"repo": "lora_repo1", "weights": "weights1.safetensors", "adapter_name": "adapter1", "adapter_weight": 1.0}]', 
+                lines=5
+            )
             image_url = gr.Text(label="Image url", placeholder="Enter image url to enable image to image model", lines=1)
             run_button = gr.Button("Run", scale=0)
 
             with gr.Accordion("Advanced Settings", open=False):
-
                 with gr.Row():
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
@@ -284,14 +280,16 @@ with gr.Blocks() as demo:
                 secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
                 bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
         
-
         with gr.Column():
+            # CHANGED: Add image output
+            output_image = gr.Image(label="Generated Image", type="pil")
             json_text = gr.Text(label="Result JSON")
             
     gr.Markdown("**Disclaimer:**")
     gr.Markdown(
         "This demo is only for research purpose. This space owner cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. This space owner provides the tools, but the responsibility for their use lies with the individual user."
     )
+    
     inputs = [
         prompt,
         image_url,
@@ -310,7 +308,8 @@ with gr.Blocks() as demo:
         bucket
     ]
 
-    outputs = [json_text]
+    # CHANGED: Two outputs now
+    outputs = [output_image, json_text]
 
     run_button.click(
         fn=run_lora,
@@ -321,4 +320,4 @@ with gr.Blocks() as demo:
 try:
     demo.queue().launch()
 except:
-   print("demo exception ...")
+    print("demo exception ...")