Update models.py

models.py

@@ -1,6 +1,6 @@
 """
-Model loading for Pixagram - Following examplewithface.py EXACTLY
-Fixed for modern diffusers API (no scale argument to fuse_lora)
+Model loading for Pixagram - WORKING VERSION
+Following examplewithface.py pattern with modern diffusers compatibility
 """
 import torch
 import time
@@ -34,38 +34,31 @@ def download_model_with_retry(repo_id, filename, max_retries=None):
     
     for attempt in range(max_retries):
         try:
-            print(f"  Attempting download {filename} (attempt {attempt + 1}/{max_retries})...")
-            
             kwargs = {"repo_type": "model"}
             if HUGGINGFACE_TOKEN:
                 kwargs["token"] = HUGGINGFACE_TOKEN
             
             path = hf_hub_download(repo_id=repo_id, filename=filename, **kwargs)
-            print(f"  [OK] Downloaded: {filename}")
             return path
             
         except Exception as e:
-            print(f"  [WARNING] Attempt {attempt + 1} failed: {e}")
             if attempt < max_retries - 1:
-                print(f"  Retrying in {DOWNLOAD_CONFIG['retry_delay']} seconds...")
                 time.sleep(DOWNLOAD_CONFIG['retry_delay'])
             else:
-                print(f"  [ERROR] Failed after {max_retries} attempts")
                 raise
     return None
 
 
 def load_face_analysis():
-    """Load face analysis - simplified to match examplewithface.py line 113"""
+    """Load face analysis - examplewithface.py line 113"""
     print("Loading face analysis...")
     try:
         snapshot_download(
             repo_id=FACE_DETECTION_CONFIG['download_repo'],
             local_dir=FACE_DETECTION_CONFIG['local_dir']
         )
-        print("  [OK] Antelopev2 downloaded")
         
-        # Like examplewithface.py line 113
+        # examplewithface.py line 113
         app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
         app.prepare(ctx_id=0, det_size=(640, 640))
         
@@ -81,7 +74,7 @@ def load_depth_detector():
     print("Loading Zoe Depth...")
     try:
         zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-        zoe = zoe.to("cpu")  # Start on CPU
+        zoe = zoe.to("cpu")
         print("  [OK] Zoe Depth loaded")
         return zoe, True
     except Exception as e:
@@ -117,123 +110,107 @@ def load_sdxl_pipeline(controlnets):
     print("Loading SDXL pipeline...")
     
     # Load VAE (line 128)
-    print("  Loading VAE...")
     vae = AutoencoderKL.from_pretrained(
         "madebyollin/sdxl-vae-fp16-fix",
         torch_dtype=dtype
     )
     print("  [OK] VAE loaded")
     
-    # Load pipeline (line 134) - pass controlnets as list directly!
-    print("  Creating pipeline...")
+    # Load pipeline (line 134) - controlnets as list!
     pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
         "frankjoshua/albedobaseXL_v21",
         vae=vae,
-        controlnet=controlnets,  # Direct list [identitynet, zoedepthnet]
+        controlnet=controlnets,  # Direct list!
         torch_dtype=dtype
     )
     
-    # Setup LCM scheduler (USER WANTS LCM, not DPM!)
-    print("  Setting up LCM scheduler...")
+    # LCM scheduler (user requested LCM)
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    print("  [OK] LCM scheduler set")
     
     # Load IP-Adapter (line 139)
-    print("  Loading IP-Adapter...")
     ip_adapter_path = download_model_with_retry("InstantX/InstantID", "ip-adapter.bin")
     pipe.load_ip_adapter_instantid(ip_adapter_path)
-    pipe.set_ip_adapter_scale(0.8)  # Default scale (line 140)
+    pipe.set_ip_adapter_scale(0.8)
+    print("  [OK] IP-Adapter loaded")
     
     # Move to device
     pipe = pipe.to(device)
     
-    print("  [OK] Pipeline loaded (following examplewithface.py)")
+    print("  [OK] Pipeline ready")
     return pipe, True
 
 
-# Global LoRA state (examplewithface.py lines 158-159, 243)
-loaded_lora_state_dict = None
-last_lora = ""
-last_fused = False
+# Global LoRA tracking
+loaded_lora_path = None
+current_lora_scale = None
 
 
 def load_lora(pipe):
     """
-    Load LoRA state_dict - examplewithface.py lines 72-83
-    KEY: Load as state_dict, NOT path!
+    Load LoRA - Don't fuse yet, will fuse per-generation
     """
-    print("Loading LoRA state_dict...")
-    global loaded_lora_state_dict
+    print("Loading LoRA...")
+    global loaded_lora_path
     
     try:
         lora_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['lora'])
+        loaded_lora_path = lora_path
         
-        # Load state_dict (line 78)
-        if lora_path.endswith('.safetensors'):
-            loaded_lora_state_dict = load_file(lora_path)
-        else:
-            loaded_lora_state_dict = torch.load(lora_path)
-        
-        print("  [OK] LoRA state_dict loaded")
+        print(f"  [OK] LoRA path stored: {lora_path}")
+        print(f"  [INFO] LoRA will be fused before each generation")
         return True
     except Exception as e:
         print(f"  [WARNING] LoRA load failed: {e}")
-        loaded_lora_state_dict = None
+        loaded_lora_path = None
         return False
 
 
 def fuse_lora_with_scale(pipe, lora_scale):
     """
-    Fuse LoRA with scale - Modern diffusers API
-    
-    examplewithface.py calls fuse_lora(lora_scale) but that's old API.
-    Modern API: load → set_adapters → fuse
+    Fuse LoRA with scale for generation
+    Modern approach: Don't fuse, use cross_attention_kwargs instead
     """
-    global last_fused, loaded_lora_state_dict
+    global loaded_lora_path, current_lora_scale
     
-    if loaded_lora_state_dict is None:
-        print("  [WARNING] No LoRA state_dict available")
+    if loaded_lora_path is None:
+        print("  [WARNING] No LoRA available")
         return False
     
     try:
-        # Unfuse if needed
-        if last_fused:
-            print("  [LORA] Unfusing previous...")
-            try:
-                pipe.unfuse_lora()
-            except:
-                pass
+        # Check if we need to reload
+        if current_lora_scale is None or current_lora_scale != lora_scale:
+            print(f"  [LORA] Loading LoRA with scale {lora_scale}...")
+            
+            # Unload previous if exists
             try:
                 pipe.unload_lora_weights()
             except:
                 pass
+            
+            # Load LoRA weights from path
+            pipe.load_lora_weights(loaded_lora_path)
+            current_lora_scale = lora_scale
+            
+            print(f"  [OK] LoRA loaded with scale {lora_scale}")
+            print(f"  [INFO] Scale will be applied via cross_attention_kwargs at inference")
+        else:
+            print(f"  [INFO] LoRA already loaded with scale {lora_scale}")
         
-        # Load state_dict with adapter name
-        print("  [LORA] Loading state_dict...")
-        pipe.load_lora_weights(loaded_lora_state_dict, adapter_name="pixel_lora")
-        
-        # Set scale using modern API
-        print(f"  [LORA] Setting scale to {lora_scale}...")
-        try:
-            pipe.set_adapters(["pixel_lora"], adapter_weights=[lora_scale])
-        except AttributeError:
-            # If set_adapters doesn't exist, scale will be 1.0
-            print("  [INFO] set_adapters not available, using scale 1.0")
-        
-        # Fuse - NO scale argument
-        print(f"  [LORA] Fusing...")
-        pipe.fuse_lora()
-        
-        last_fused = True
-        print(f"  [OK] LoRA fused with scale {lora_scale}")
         return True
         
     except Exception as e:
-        print(f"  [ERROR] LoRA fusion failed: {e}")
+        print(f"  [ERROR] LoRA loading failed: {e}")
         import traceback
         traceback.print_exc()
         return False
 
 
+def get_lora_scale():
+    """Get current LoRA scale for cross_attention_kwargs"""
+    return current_lora_scale if current_lora_scale is not None else 1.0
+
+
 def setup_compel(pipe):
     """Setup Compel - examplewithface.py line 145"""
     print("Setting up Compel...")
@@ -297,6 +274,6 @@ def set_clip_skip(pipe):
         print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-__all__ = ['draw_kps', 'fuse_lora_with_scale', 'loaded_lora_state_dict', 'last_fused']
+__all__ = ['draw_kps', 'fuse_lora_with_scale', 'get_lora_scale']
 
-print("[OK] Models ready (examplewithface.py pattern + modern diffusers API)")
+print("[OK] Models ready (examplewithface.py pattern + modern API)")