Update models.py

models.py

@@ -1,6 +1,6 @@
 """
-Model loading and initialization for Pixagram AI Pixel Art Generator
-FIXED VERSION - Following examplewithface.py EXACTLY for LORA handling
+Model loading for Pixagram - Following examplewithface.py EXACTLY
+Fixed for modern diffusers API (no scale argument to fuse_lora)
 """
 import torch
 import time
@@ -13,10 +13,9 @@ from diffusers import (
 from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector
 from huggingface_hub import hf_hub_download, snapshot_download
-from safetensors.torch import load_file  # CRITICAL: For loading state_dict
+from safetensors.torch import load_file
 from compel import Compel, ReturnedEmbeddingsType
 
-# Use InstantID pipeline
 from pipeline_stable_diffusion_xl_instantid_img2img import (
     StableDiffusionXLInstantIDImg2ImgPipeline,
     draw_kps
@@ -29,94 +28,71 @@ from config import (
 
 
 def download_model_with_retry(repo_id, filename, max_retries=None):
-    """Download model with retry logic and proper token handling."""
+    """Download model with retry logic"""
     if max_retries is None:
         max_retries = DOWNLOAD_CONFIG['max_retries']
     
     for attempt in range(max_retries):
         try:
-            print(f"  Attempting to download {filename} (attempt {attempt + 1}/{max_retries})...")
+            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
-            )
+            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] Download attempt {attempt + 1} failed: {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 to download {filename} after {max_retries} attempts")
+                print(f"  [ERROR] Failed after {max_retries} attempts")
                 raise
-    
     return None
 
 
 def load_face_analysis():
-    """
-    Load face analysis model - SIMPLIFIED to match working example
-    """
-    print("Loading face analysis model...")
-    
+    """Load face analysis - simplified to match examplewithface.py line 113"""
+    print("Loading face analysis...")
     try:
-        # Download antelopev2 model files
-        print("  Downloading antelopev2 model files...")
         snapshot_download(
             repo_id=FACE_DETECTION_CONFIG['download_repo'],
             local_dir=FACE_DETECTION_CONFIG['local_dir']
         )
         print("  [OK] Antelopev2 downloaded")
         
-        # Initialize FaceAnalysis (like examplewithface.py line 113)
-        face_app = FaceAnalysis(
-            name='antelopev2',
-            root='/data',
-            providers=['CPUExecutionProvider']
-        )
-        
-        # Prepare the model (like examplewithface.py line 114)
-        face_app.prepare(ctx_id=0, det_size=(640, 640))
-        
-        print(f"  [OK] Face analysis model loaded successfully")
-        return face_app, True
+        # Like examplewithface.py line 113
+        app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
+        app.prepare(ctx_id=0, det_size=(640, 640))
         
+        print("  [OK] Face analysis loaded")
+        return app, True
     except Exception as e:
-        print(f"  [ERROR] Face analysis loading failed: {e}")
-        import traceback
-        traceback.print_exc()
+        print(f"  [ERROR] Face analysis failed: {e}")
         return None, False
 
 
 def load_depth_detector():
-    """Load Zoe Depth detector (like examplewithface.py line 151)"""
-    print("Loading Zoe Depth detector...")
+    """Load Zoe Depth - examplewithface.py line 151"""
+    print("Loading Zoe Depth...")
     try:
-        zoe_depth = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-        # Start on CPU to save memory
-        zoe_depth = zoe_depth.to("cpu")
-        print("  [OK] Zoe Depth loaded (on CPU, will move to GPU when needed)")
-        return zoe_depth, True
+        zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
+        zoe = zoe.to("cpu")  # Start on CPU
+        print("  [OK] Zoe Depth loaded")
+        return zoe, True
     except Exception as e:
-        print(f"  [WARNING] Zoe Depth not available: {e}")
+        print(f"  [WARNING] Zoe Depth unavailable: {e}")
         return None, False
 
 
 def load_controlnets():
-    """
-    Load ControlNets for InstantID pipeline
-    Following examplewithface.py lines 122-126
-    """
-    print("Loading InstantID ControlNet...")
+    """Load ControlNets - examplewithface.py lines 122-126"""
+    print("Loading ControlNets...")
+    
     identitynet = ControlNetModel.from_pretrained(
         "InstantX/InstantID",
         subfolder="ControlNetModel",
@@ -124,7 +100,6 @@ def load_controlnets():
     )
     print("  [OK] InstantID ControlNet loaded")
     
-    print("Loading Zoe Depth ControlNet...")
     zoedepthnet = ControlNetModel.from_pretrained(
         "diffusers/controlnet-zoe-depth-sdxl-1.0",
         torch_dtype=dtype
@@ -136,133 +111,132 @@ def load_controlnets():
 
 def load_sdxl_pipeline(controlnets):
     """
-    Load SDXL pipeline with InstantID support
-    Following examplewithface.py lines 128-145 EXACTLY
+    Load pipeline - examplewithface.py lines 128-145
+    KEY: Pass controlnets as LIST directly, NO wrapper
     """
-    print("Loading SDXL checkpoint with InstantID pipeline...")
+    print("Loading SDXL pipeline...")
     
-    try:
-        # Load VAE first (line 128)
-        print("Loading FP16-fixed VAE...")
-        vae = AutoencoderKL.from_pretrained(
-            "madebyollin/sdxl-vae-fp16-fix",
-            torch_dtype=dtype
-        )
-        print("  [OK] VAE loaded")
-        
-        # Create pipeline with direct controlnet list (line 134)
-        print("Creating InstantID pipeline...")
-        pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
-            "frankjoshua/albedobaseXL_v21",
-            vae=vae,
-            controlnet=controlnets,  # Pass list directly [identitynet, zoedepthnet]
-            torch_dtype=dtype
-        )
-        
-        # Use LCM scheduler (user wants LCM, not DPM)
-        print("Setting up LCM scheduler...")
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-        print("  [OK] LCM scheduler configured")
-        
-        # Load IP-Adapter (line 139)
-        print("Loading IP-Adapter for InstantID...")
-        ip_adapter_path = download_model_with_retry(
-            "InstantX/InstantID",
-            "ip-adapter.bin"
-        )
-        pipe.load_ip_adapter_instantid(ip_adapter_path)
-        
-        # Set initial scale (line 140)
-        pipe.set_ip_adapter_scale(0.8)
-        print("  [OK] IP-Adapter loaded with scale 0.8")
-        
-        # Move to device
-        pipe = pipe.to(device)
-        
-        print("  [OK] InstantID pipeline loaded (following examplewithface.py)")
-        return pipe, True
-        
-    except Exception as e:
-        print(f"  [ERROR] Could not load InstantID pipeline: {e}")
-        import traceback
-        traceback.print_exc()
-        raise
+    # 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...")
+    pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
+        "frankjoshua/albedobaseXL_v21",
+        vae=vae,
+        controlnet=controlnets,  # Direct list [identitynet, zoedepthnet]
+        torch_dtype=dtype
+    )
+    
+    # Setup LCM scheduler (USER WANTS LCM, not DPM!)
+    print("  Setting up LCM scheduler...")
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    
+    # 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)
+    
+    # Move to device
+    pipe = pipe.to(device)
+    
+    print("  [OK] Pipeline loaded (following examplewithface.py)")
+    return pipe, True
 
 
-# Global variable to track LORA state (like examplewithface.py lines 158-159)
-last_lora_fused = False
+# Global LoRA state (examplewithface.py lines 158-159, 243)
 loaded_lora_state_dict = None
+last_lora = ""
+last_fused = False
 
 
 def load_lora(pipe):
     """
-    Load LORA state_dict into memory
-    Following examplewithface.py lines 72-83 pattern
+    Load LoRA state_dict - examplewithface.py lines 72-83
+    KEY: Load as state_dict, NOT path!
     """
-    print("Loading LORA state_dict from HuggingFace Hub...")
+    print("Loading LoRA state_dict...")
     global loaded_lora_state_dict
     
     try:
         lora_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['lora'])
         
-        # Load state_dict like the working example (line 78)
+        # 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(f"  [OK] LORA state_dict loaded (will be fused before generation)")
+        print("  [OK] LoRA state_dict loaded")
         return True
-        
     except Exception as e:
-        print(f"  [WARNING] Could not load LORA: {e}")
+        print(f"  [WARNING] LoRA load failed: {e}")
         loaded_lora_state_dict = None
         return False
 
 
 def fuse_lora_with_scale(pipe, lora_scale):
     """
-    Fuse LORA with scale before generation
-    Following examplewithface.py lines 256-271 pattern EXACTLY
+    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
     """
-    global last_lora_fused, loaded_lora_state_dict
+    global last_fused, loaded_lora_state_dict
     
     if loaded_lora_state_dict is None:
-        print("  [WARNING] No LORA state_dict loaded")
+        print("  [WARNING] No LoRA state_dict available")
         return False
     
     try:
-        # Unfuse if already fused (line 257-264)
-        if last_lora_fused:
-            print(f"  [LORA] Unfusing previous...")
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
+        # Unfuse if needed
+        if last_fused:
+            print("  [LORA] Unfusing previous...")
+            try:
+                pipe.unfuse_lora()
+            except:
+                pass
+            try:
+                pipe.unload_lora_weights()
+            except:
+                pass
         
-        # Load and fuse with scale (lines 266-267)
-        print(f"  [LORA] Loading state_dict...")
-        pipe.load_lora_weights(loaded_lora_state_dict)
+        # Load state_dict with adapter name
+        print("  [LORA] Loading state_dict...")
+        pipe.load_lora_weights(loaded_lora_state_dict, adapter_name="pixel_lora")
         
-        print(f"  [LORA] Fusing with scale: {lora_scale}")
-        pipe.fuse_lora(lora_scale)
+        # 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")
         
-        last_lora_fused = True
-        print(f"  [OK] LORA fused successfully")
+        # 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"  [WARNING] Could not fuse LORA: {e}")
+        print(f"  [ERROR] LoRA fusion failed: {e}")
         import traceback
         traceback.print_exc()
         return False
 
 
 def setup_compel(pipe):
-    """
-    Setup Compel for SDXL prompt handling
-    Following examplewithface.py line 145 pattern
-    """
-    print("Setting up Compel for enhanced prompt processing...")
+    """Setup Compel - examplewithface.py line 145"""
+    print("Setting up Compel...")
     try:
         compel = Compel(
             tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
@@ -270,78 +244,59 @@ def setup_compel(pipe):
             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
             requires_pooled=[False, True]
         )
-        print("  [OK] Compel loaded successfully")
+        print("  [OK] Compel loaded")
         return compel, True
     except Exception as e:
-        print(f"  [WARNING] Compel not available: {e}")
+        print(f"  [WARNING] Compel unavailable: {e}")
         return None, False
 
 
 def setup_scheduler(pipe):
-    """Setup scheduler - already done in load_sdxl_pipeline()"""
+    """Already done in load_sdxl_pipeline"""
     pass
 
 
 def optimize_pipeline(pipe):
-    """Apply optimizations to pipeline"""
+    """Apply optimizations"""
     if device == "cuda":
         try:
             pipe.enable_xformers_memory_efficient_attention()
             print("  [OK] xformers enabled")
-        except Exception as e:
-            print(f"  [INFO] xformers not available: {e}")
+        except:
+            pass
     
-    # VAE optimizations
     if hasattr(pipe, 'enable_vae_slicing'):
         pipe.enable_vae_slicing()
-        print("  [OK] VAE slicing enabled")
-    
     if hasattr(pipe, 'enable_vae_tiling'):
         pipe.enable_vae_tiling()
-        print("  [OK] VAE tiling enabled")
 
 
 def load_caption_model():
-    """Load caption model with proper error handling"""
+    """Load caption model"""
     print("Loading caption model...")
-    
     try:
         from transformers import AutoProcessor, AutoModelForCausalLM
-        
-        print("  Attempting GIT-Large...")
-        caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
-        caption_model = AutoModelForCausalLM.from_pretrained(
-            "microsoft/git-large-coco",
-            torch_dtype=dtype
-        ).to("cpu")
-        print("  [OK] GIT-Large model loaded")
-        return caption_processor, caption_model, True, 'git'
-    except Exception as e1:
-        print(f"  [INFO] GIT-Large not available: {e1}")
-        
+        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")
+        return processor, model, True, 'git'
+    except:
         try:
             from transformers import BlipProcessor, BlipForConditionalGeneration
-            
-            print("  Attempting BLIP base...")
-            caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-            caption_model = BlipForConditionalGeneration.from_pretrained(
-                "Salesforce/blip-image-captioning-base",
-                torch_dtype=dtype
-            ).to("cpu")
-            print("  [OK] BLIP base model loaded")
-            return caption_processor, caption_model, True, 'blip'
-        except Exception as e2:
-            print(f"  [WARNING] Caption models not available: {e2}")
+            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")
+            return processor, model, True, 'blip'
+        except:
             return None, None, False, 'none'
 
 
 def set_clip_skip(pipe):
-    """Set CLIP skip value"""
+    """Set CLIP skip"""
     if hasattr(pipe, 'text_encoder'):
         print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-# Export for use in generator
-__all__ = ['draw_kps', 'fuse_lora_with_scale', 'loaded_lora_state_dict', 'last_lora_fused']
+__all__ = ['draw_kps', 'fuse_lora_with_scale', 'loaded_lora_state_dict', 'last_fused']
 
-print("[OK] Model loading functions ready (following examplewithface.py EXACTLY)")
+print("[OK] Models ready (examplewithface.py pattern + modern diffusers API)")