Update models.py

models.py

@@ -1,15 +1,12 @@
 """
-Model loading for Pixagram - WORKING VERSION
-Following examplewithface.py pattern with modern diffusers compatibility
+Models.py - Following examplewithface.py EXACTLY
+NO MultiControlNetModel wrapper!
+NO fuse_lora with scale!
 """
 import torch
 import time
 import os
-from diffusers import (
-    ControlNetModel,
-    AutoencoderKL,
-    LCMScheduler
-)
+from diffusers import ControlNetModel, AutoencoderKL, LCMScheduler
 from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -28,7 +25,6 @@ from config import (
 
 
 def download_model_with_retry(repo_id, filename, max_retries=None):
-    """Download model with retry logic"""
     if max_retries is None:
         max_retries = DOWNLOAD_CONFIG['max_retries']
     
@@ -40,7 +36,6 @@ def download_model_with_retry(repo_id, filename, max_retries=None):
             
             path = hf_hub_download(repo_id=repo_id, filename=filename, **kwargs)
             return path
-            
         except Exception as e:
             if attempt < max_retries - 1:
                 time.sleep(DOWNLOAD_CONFIG['retry_delay'])
@@ -50,7 +45,7 @@ def download_model_with_retry(repo_id, filename, max_retries=None):
 
 
 def load_face_analysis():
-    """Load face analysis - examplewithface.py line 113"""
+    """examplewithface.py line 113"""
     print("Loading face analysis...")
     try:
         snapshot_download(
@@ -58,7 +53,6 @@ def load_face_analysis():
             local_dir=FACE_DETECTION_CONFIG['local_dir']
         )
         
-        # examplewithface.py line 113
         app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
         app.prepare(ctx_id=0, det_size=(640, 640))
         
@@ -70,20 +64,19 @@ def load_face_analysis():
 
 
 def load_depth_detector():
-    """Load Zoe Depth - examplewithface.py line 151"""
+    """examplewithface.py line 151"""
     print("Loading Zoe Depth...")
     try:
         zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-        zoe = zoe.to("cpu")
         print("  [OK] Zoe Depth loaded")
         return zoe, True
     except Exception as e:
-        print(f"  [WARNING] Zoe Depth unavailable: {e}")
+        print(f"  [WARNING] Zoe unavailable: {e}")
         return None, False
 
 
 def load_controlnets():
-    """Load ControlNets - examplewithface.py lines 122-126"""
+    """examplewithface.py lines 122-126"""
     print("Loading ControlNets...")
     
     identitynet = ControlNetModel.from_pretrained(
@@ -91,23 +84,23 @@ def load_controlnets():
         subfolder="ControlNetModel",
         torch_dtype=dtype
     )
-    print("  [OK] InstantID ControlNet loaded")
+    print("  [OK] InstantID ControlNet")
     
     zoedepthnet = ControlNetModel.from_pretrained(
         "diffusers/controlnet-zoe-depth-sdxl-1.0",
         torch_dtype=dtype
     )
-    print("  [OK] Zoe Depth ControlNet loaded")
+    print("  [OK] Zoe Depth ControlNet")
     
     return identitynet, zoedepthnet
 
 
 def load_sdxl_pipeline(controlnets):
     """
-    Load pipeline - examplewithface.py lines 128-145
-    KEY: Pass controlnets as LIST directly, NO wrapper
+    examplewithface.py lines 128-145
+    CRITICAL: Pass controlnets as LIST - NO MultiControlNetModel!
     """
-    print("Loading SDXL pipeline...")
+    print("Loading pipeline...")
     
     # Load VAE (line 128)
     vae = AutoencoderKL.from_pretrained(
@@ -116,103 +109,78 @@ def load_sdxl_pipeline(controlnets):
     )
     print("  [OK] VAE loaded")
     
-    # Load pipeline (line 134) - controlnets as list!
+    # Create pipeline (line 134) - controlnets as LIST!
     pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
         "frankjoshua/albedobaseXL_v21",
         vae=vae,
-        controlnet=controlnets,  # Direct list!
+        controlnet=controlnets,  # ← LIST [identitynet, zoedepthnet] - NO WRAPPER!
         torch_dtype=dtype
     )
+    print("  [OK] Pipeline created with direct controlnet list")
     
-    # LCM scheduler (user requested LCM)
+    # LCM scheduler
     pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-    print("  [OK] LCM scheduler set")
+    print("  [OK] LCM scheduler")
     
-    # Load IP-Adapter (line 139)
+    # IP-Adapter (line 139)
     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)
     print("  [OK] IP-Adapter loaded")
     
-    # Move to device
     pipe = pipe.to(device)
-    
-    print("  [OK] Pipeline ready")
+    print("  [OK] Pipeline ready (following examplewithface.py EXACTLY)")
     return pipe, True
 
 
-# Global LoRA tracking
-loaded_lora_path = None
-current_lora_scale = None
+# Global LoRA state
+lora_path_cached = None
 
 
 def load_lora(pipe):
-    """
-    Load LoRA - Don't fuse yet, will fuse per-generation
-    """
+    """Load LoRA - store path for later use"""
     print("Loading LoRA...")
-    global loaded_lora_path
+    global lora_path_cached
     
     try:
         lora_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['lora'])
-        loaded_lora_path = lora_path
-        
-        print(f"  [OK] LoRA path stored: {lora_path}")
-        print(f"  [INFO] LoRA will be fused before each generation")
+        lora_path_cached = lora_path
+        print(f"  [OK] LoRA path stored")
         return True
     except Exception as e:
-        print(f"  [WARNING] LoRA load failed: {e}")
-        loaded_lora_path = None
+        print(f"  [WARNING] LoRA failed: {e}")
         return False
 
 
 def fuse_lora_with_scale(pipe, lora_scale):
     """
-    Fuse LoRA with scale for generation
-    Modern approach: Don't fuse, use cross_attention_kwargs instead
+    Modern approach: Load LoRA and let cross_attention_kwargs apply scale
     """
-    global loaded_lora_path, current_lora_scale
+    global lora_path_cached
     
-    if loaded_lora_path is None:
-        print("  [WARNING] No LoRA available")
+    if lora_path_cached is None:
         return False
     
     try:
-        # 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}")
+        # Unload previous
+        try:
+            pipe.unload_lora_weights()
+        except:
+            pass
         
-        return True
+        # Load LoRA
+        print(f"  [LORA] Loading with scale {lora_scale}...")
+        pipe.load_lora_weights(lora_path_cached)
+        print(f"  [OK] LoRA loaded (scale will be applied via cross_attention_kwargs)")
         
+        return True
     except Exception as e:
-        print(f"  [ERROR] LoRA loading failed: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"  [ERROR] LoRA failed: {e}")
         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"""
+    """examplewithface.py line 145"""
     print("Setting up Compel...")
     try:
         compel = Compel(
@@ -221,7 +189,7 @@ def setup_compel(pipe):
             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
             requires_pooled=[False, True]
         )
-        print("  [OK] Compel loaded")
+        print("  [OK] Compel ready")
         return compel, True
     except Exception as e:
         print(f"  [WARNING] Compel unavailable: {e}")
@@ -229,12 +197,10 @@ def setup_compel(pipe):
 
 
 def setup_scheduler(pipe):
-    """Already done in load_sdxl_pipeline"""
     pass
 
 
 def optimize_pipeline(pipe):
-    """Apply optimizations"""
     if device == "cuda":
         try:
             pipe.enable_xformers_memory_efficient_attention()
@@ -249,31 +215,29 @@ def optimize_pipeline(pipe):
 
 
 def load_caption_model():
-    """Load caption model"""
     print("Loading caption model...")
     try:
         from transformers import AutoProcessor, AutoModelForCausalLM
         processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
         model = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco", torch_dtype=dtype).to("cpu")
-        print("  [OK] GIT-Large loaded")
+        print("  [OK] GIT-Large")
         return processor, model, True, 'git'
     except:
         try:
             from transformers import BlipProcessor, BlipForConditionalGeneration
             processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
             model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=dtype).to("cpu")
-            print("  [OK] BLIP loaded")
+            print("  [OK] BLIP")
             return processor, model, True, 'blip'
         except:
             return None, None, False, 'none'
 
 
 def set_clip_skip(pipe):
-    """Set CLIP skip"""
     if hasattr(pipe, 'text_encoder'):
-        print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
+        print(f"  [OK] CLIP skip {CLIP_SKIP}")
 
 
-__all__ = ['draw_kps', 'fuse_lora_with_scale', 'get_lora_scale']
+__all__ = ['draw_kps', 'fuse_lora_with_scale']
 
-print("[OK] Models ready (examplewithface.py pattern + modern API)")
+print("[OK] models.py ready - NO MultiControlNetModel, following examplewithface.py")