Update utils.py

utils.py

@@ -19,7 +19,6 @@ import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
 
-
 from src.flux.generate import generate, seed_everything
 
 try:
@@ -30,54 +29,82 @@ except ImportError:
 
 import re
 
+# Global variables
 pipe = None
 model_dict = {}
-
+_MODEL_INITIALIZED = False
+_ADAPTERS_LOADED = False
 
 def init_flux_pipeline():
-    global pipe
+    """Initialize Flux model, ensuring it runs only once"""
+    global pipe, _MODEL_INITIALIZED
+    
     if pipe is None:
+        print("Initializing Flux pipeline...")
         token = os.getenv("HF_TOKEN")
         if not token:
             raise ValueError("HF_TOKEN environment variable not set.")
+            
         pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-schnell",
             use_auth_token=token,
             torch_dtype=torch.bfloat16
         )
         pipe = pipe.to("cuda")
+        _MODEL_INITIALIZED = True
+        print("Flux pipeline initialized successfully.")
+    
+    return pipe
 
 def load_all_lora_adapters():
-    global pipe
-    if pipe is None:
-        init_flux_pipeline()
-
-    # Define all LoRA adapters to preload
-    LORA_ADAPTERS = {
-        "add": "weights/add.safetensors",
-        "remove": "weights/remove.safetensors",
-        "action": "weights/action.safetensors",
-        "expression": "weights/expression.safetensors",
-        "addition": "weights/addition.safetensors",
-        "material": "weights/material.safetensors",
-        "color": "weights/color.safetensors",
-        "bg": "weights/bg.safetensors",
-        "appearance": "weights/appearance.safetensors",
-        "fusion": "weights/fusion.safetensors",
-        "overall": "weights/overall.safetensors",
-    }
-
-    print(LORA_ADAPTERS)
-    
-    for adapter_name, weight_path in LORA_ADAPTERS.items():
-        pipe.load_lora_weights(
-            "Cicici1109/IEAP",
-            weight_name=weight_path,
-            adapter_name=adapter_name,
-        )
-        print(f"Loaded LoRA adapter: {adapter_name}")
+    """Load all LoRA adapters, ensuring it runs only once"""
+    global pipe, _ADAPTERS_LOADED
+    
+    # Ensure model is initialized
+    init_flux_pipeline()
+    
+    if not _ADAPTERS_LOADED:
+        print("Loading all LoRA adapters...")
+        
+        LORA_ADAPTERS = {
+            "add": "weights/add.safetensors",
+            "remove": "weights/remove.safetensors",
+            "action": "weights/action.safetensors",
+            "expression": "weights/expression.safetensors",
+            "addition": "weights/addition.safetensors",
+            "material": "weights/material.safetensors",
+            "color": "weights/color.safetensors",
+            "bg": "weights/bg.safetensors",
+            "appearance": "weights/appearance.safetensors",
+            "fusion": "weights/fusion.safetensors",
+            "overall": "weights/overall.safetensors",
+        }
+        
+        for adapter_name, weight_path in LORA_ADAPTERS.items():
+            if not os.path.exists(weight_path):
+                print(f"Warning: LoRA weight file not found: {weight_path}")
+                continue
+                
+            try:
+                pipe.load_lora_weights(
+                    "Cicici1109/IEAP",
+                    weight_name=weight_path,
+                    adapter_name=adapter_name,
+                )
+                print(f"✅ Successfully loaded adapter: {adapter_name}")
+            except Exception as e:
+                print(f"❌ Failed to load adapter {adapter_name}: {e}")
+        
+        loaded_adapters = list(pipe.lora_adapters.keys())
+        print(f"Loaded adapters: {loaded_adapters}")
+        
+        if loaded_adapters:
+            pipe.set_adapters(loaded_adapters[0])
+            print(f"Default adapter set to: {loaded_adapters[0]}")
+        
+        _ADAPTERS_LOADED = True
     
-    pipe.set_adapters("scene")  # Set default adapter
+    return pipe
 
 def get_model(model_path):
     global model_dict
@@ -198,8 +225,9 @@ def extract_last_bbox(result):
 
 @spaces.GPU
 def infer_with_DiT(task, image, instruction, category):
-    init_flux_pipeline()
-
+    # Ensure model and adapters are initialized
+    load_all_lora_adapters()
+    
     if task == 'RoI Inpainting':
         if category == 'Add' or category == 'Replace':
             adapter_name = "add"
@@ -243,7 +271,8 @@ def infer_with_DiT(task, image, instruction, category):
     else:
         raise ValueError(f"Invalid task: '{task}'")
     
-    # Switch to the preloaded adapter
+    # Switch to the specified adapter
+    print(f"Switching to adapter: {adapter_name}")
     pipe.set_adapters(adapter_name)
     
     result_img = generate(
@@ -621,8 +650,4 @@ def layout_change(bbox, instruction):
     result = response.choices[0].message.content.strip()
 
     bbox = extract_last_bbox(result)
-    return bbox
-
-if __name__ == "__main__":
-    init_flux_pipeline()
-    load_all_lora_adapters()  # Preload all LoRA adapters at startup
+    return bbox