Update app.py

app.py

@@ -38,13 +38,13 @@ 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"],
-)
+# 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, quantization_config=pipeline_quant_config, torch_dtype=dtype)
+txt2img_pipe = FluxKontextPipeline.from_pretrained(base_model, torch_dtype=dtype)
 txt2img_pipe = txt2img_pipe.to(device)
 
 MAX_SEED = 2**32 - 1
@@ -69,13 +69,16 @@ 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"):
-        print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
         connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
-    
         s3 = boto3.client(
             's3',
             endpoint_url=connectionUrl,
@@ -113,97 +116,116 @@ def generate_random_4_digit_string():
     return ''.join(random.choices(string.digits, k=4))
 
 @spaces.GPU(duration=120)
-def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
+def run_lora(
+    prompt, 
+    image_url, 
+    lora_strings_json, 
+    image_strength,  
+    cfg_scale, 
+    steps, 
+    randomize_seed, 
+    seed, 
+    width, 
+    height, 
+    upload_to_r2, 
+    account_id, 
+    access_key, 
+    secret_key, 
+    bucket, 
+    progress=gr.Progress(track_tqdm=True)
+):
     print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height)
     gr.Info("Starting process")
-    img2img_model = False
-    orginal_image = None
-    device = txt2img_pipe.device
-    
+    pipe = txt2img_pipe
+    device = pipe.device
     print(device)
-  
-    # Set random seed for reproducibility
+
+    # ========== Seed ==========
     if randomize_seed:
         with calculateDuration("Set random seed"):
             seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Load LoRA weights
+    # ========== LoRA ==========
     gr.Info("Start to load LoRA ...")
     with calculateDuration("Unloading LoRA"):
-        txt2img_pipe.unload_lora_weights()
-    
-    lora_configs = None
+        try:
+            pipe.unload_lora_weights()
+        except Exception as _:
+            # 某些版本上未加载时调用可能抛异常，忽略
+            pass
+
     adapter_names = []
-    lora_names = []
+    adapter_weights = []
     if lora_strings_json:
         try:
             lora_configs = json.loads(lora_strings_json)
-        except:
+        except Exception as _:
+            lora_configs = None
             gr.Warning("Parse lora config json failed")
             print("parse lora config json failed")
-            
+
         if lora_configs:
-           
             with calculateDuration("Loading LoRA weights"):
-                adapter_weights = []
-                
-                for idx, lora_info in enumerate(lora_configs):
-                    lora_repo = lora_info.get("repo")
+                for lora_info in lora_configs:
+                    repo = lora_info.get("repo")
                     weights = lora_info.get("weights")
-                    adapter_name = lora_info.get("adapter_name")
-                    
-                    lora_name = generate_random_4_digit_string()
-                    lora_names.append(lora_name)
-                    adapter_weight = lora_info.get("adapter_weight")
-                    adapter_names.append(adapter_name)
-                    adapter_weights.append(adapter_weight)
-                    if lora_repo and weights and adapter_name:
-                        try:
-                            txt2img_pipe.load_lora_weights(hf_hub_download(lora_repo, weights), adapter_name=lora_name)
-                        except:
-                            print("load lora error")
-                
-                # set lora weights
-                if len(lora_names) > 0:
-                    txt2img_pipe.set_adapters(lora_names, adapter_weights=adapter_weights)
-                    txt2img_pipe.fuse_lora(adapter_names=lora_names)
-                    txt2img_pipe.enable_vae_slicing()
-    
-    # Generate image
+                    # 优先使用用户提供的 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):
+                        print(f"skip invalid lora entry: {lora_info}")
+                        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)
+                    except Exception as e:
+                        print(f"load lora error for {repo}/{weights}: {e}")
+
+                if adapter_names:
+                    pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+                    # 可选：融合后推理更快，但无法动态调整权重
+                    pipe.fuse_lora(adapter_names=adapter_names)
+
+    pipe.enable_vae_slicing()
+
+ 
+    clip_side_prompt = safe_trim_for_clip(prompt, max_words=77) 
+    init_image = None
     error_message = ""
     try:
         gr.Info("Start to generate images ...")
-        print(device)
-        # Generate image
-        pipe = txt2img_pipe.to(device)
-        generator = torch.Generator("cuda").manual_seed(seed)
         joint_attention_kwargs = {"scale": 1}
-        final_image = pipe(
+        image = pipe(
             prompt=prompt,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
+            num_inference_steps=int(steps),
+            guidance_scale=float(cfg_scale),
+            width=int(width),
+            height=int(height),
             max_sequence_length=512,
             generator=generator,
             joint_attention_kwargs=joint_attention_kwargs
         ).images[0]
     except Exception as e:
-        error_message =  str(e)
+        error_message = str(e)
         gr.Error(error_message)
         print("fatal error", e)
-        final_image = None
-        
-    if final_image:
-        if upload_to_r2:
-            url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
-            result = {"status": "success", "message": "upload image success", "url": url}    
-        else:
-            result = {"status": "success", "message": "Image generated but not uploaded"}
-    else:
-        result = {"status": "failed", "message": error_message}
-        final_image = None
-        
+        image = None
+
+    result = {"status": "failed", "message": error_message} if image is None else {"status": "success", "message": "Image generated but not uploaded"}
+    if image is not None and upload_to_r2:
+        try:
+            url = upload_image_to_r2(image, account_id, access_key, secret_key, bucket)
+            result = {"status": "success", "message": "upload image success", "url": url}
+        except Exception as e:
+            err = f"Upload failed: {e}"
+            gr.Warning(err)
+            print(err)
+            result = {"status": "success", "message": "generated but upload failed"}
+
     gr.Info("Completed！")
     progress(100, "Completed!")
     return json.dumps(result)