Update utils.py

utils.py

@@ -19,6 +19,7 @@ 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:
@@ -29,78 +30,21 @@ except ImportError:
 
 import re
 
-# Global variables
 pipe = None
 model_dict = {}
-_MODEL_INITIALIZED = False
-_ADAPTERS_LOADED = False
 
 def init_flux_pipeline():
-    """Initialize Flux model, ensuring it runs only once"""
-    global pipe, _MODEL_INITIALIZED
-    
+    global pipe
     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():
-    """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():
-            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
-    
-    return pipe
 
 def get_model(model_path):
     global model_dict
@@ -221,55 +165,57 @@ def extract_last_bbox(result):
 
 @spaces.GPU
 def infer_with_DiT(task, image, instruction, category):
-    # Ensure model and adapters are initialized
-    load_all_lora_adapters()
-    
+    init_flux_pipeline()
+
     if task == 'RoI Inpainting':
         if category == 'Add' or category == 'Replace':
-            adapter_name = "add"
+            lora_path = "weights/add.safetensors"
             added = extract_object_with_gpt(instruction)
             instruction_dit = f"add {added} on the black region"
         elif category == 'Remove' or category == 'Action Change':
-            adapter_name = "remove"
+            lora_path = "weights/remove.safetensors"
             instruction_dit = f"Fill the hole of the image"
+       
         condition = Condition("scene", image, position_delta=(0, 0))
-    
     elif task == 'RoI Editing':
         image = Image.open(image).convert('RGB').resize((512, 512))
         condition = Condition("scene", image, position_delta=(0, -32))
         instruction_dit = instruction
-        
         if category == 'Action Change':
-            adapter_name = "action"
+            lora_path = "weights/action.safetensors"
         elif category == 'Expression Change':
-            adapter_name = "expression"
+            lora_path = "weights/expression.safetensors"
         elif category == 'Add':
-            adapter_name = "addition"
+            lora_path = "weights/addition.safetensors"
         elif category == 'Material Change':
-            adapter_name = "material"
+            lora_path = "weights/material.safetensors"
         elif category == 'Color Change':
-            adapter_name = "color"
+            lora_path = "weights/color.safetensors"
         elif category == 'Background Change':
-            adapter_name = "bg"
+            lora_path = "weights/bg.safetensors"
         elif category == 'Appearance Change':
-            adapter_name = "appearance"
-    
+            lora_path = "weights/appearance.safetensors"
+        
     elif task == 'RoI Compositioning':
-        adapter_name = "fusion"
+        lora_path = "weights/fusion.safetensors"
         condition = Condition("scene", image, position_delta=(0, 0))
         instruction_dit = "inpaint the black-bordered region so that the object's edges blend smoothly with the background"
 
     elif task == 'Global Transformation':
         image = Image.open(image).convert('RGB').resize((512, 512))
         instruction_dit = instruction
-        adapter_name = "overall"
+        lora_path = "weights/overall.safetensors"
+
         condition = Condition("scene", image, position_delta=(0, -32))
     else:
         raise ValueError(f"Invalid task: '{task}'")
     
-    # Switch to the specified adapter
-    print(f"Switching to adapter: {adapter_name}")
-    pipe.set_adapters(adapter_name)
+    pipe.unload_lora_weights()
+    pipe.load_lora_weights(
+        "Cicici1109/IEAP",
+        weight_name=lora_path,
+        adapter_name="scene",
+    )
     
     result_img = generate(
         pipe,
@@ -646,4 +592,7 @@ def layout_change(bbox, instruction):
     result = response.choices[0].message.content.strip()
 
     bbox = extract_last_bbox(result)
-    return bbox
+    return bbox
+
+if __name__ == "__main__":
+    init_flux_pipeline()