Upload 3 files

generator.py

@@ -1,6 +1,6 @@
 """
 Generation logic for Pixagram AI Pixel Art Generator  
-CLEANED VERSION - Applied all audit fixes
+CORRECTED VERSION - Following examplewithface.py pattern
 """
 import torch
 import numpy as np
@@ -21,6 +21,7 @@ from models import (
     load_sdxl_pipeline, load_lora, setup_compel,
     setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip
 )
+from pipeline_stable_diffusion_xl_instantid_img2img import draw_kps
 from memory_utils import MemoryManager, ModelOffloader
 
 
@@ -345,8 +346,7 @@ class RetroArtConverter:
                         # Get raw embeddings
                         face_embeddings = face.normed_embedding
                         
-                        # Draw keypoints
-                        from pipeline_stable_diffusion_xl_instantid_img2img import draw_kps
+                        # Draw keypoints using imported draw_kps
                         face_kps_image = draw_kps(resized_image, face.kps)
                         
                         # Get bbox for color matching
@@ -372,13 +372,14 @@ class RetroArtConverter:
             elif not self.face_detection_enabled:
                 print("[FACE] Face detection disabled - generating without face preservation")
             
-            # Set LORA scale
-            if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
+            # Fuse LORA with scale (following working example approach)
+            if self.models_loaded['lora']:
                 try:
-                    self.pipe.set_adapters(["retroart"], adapter_weights=[lora_scale])
-                    print(f"[LORA] Scale: {lora_scale}")
+                    from models import fuse_lora_with_scale
+                    fuse_lora_with_scale(self.pipe, lora_scale)
+                    print(f"[LORA] Fused with scale: {lora_scale}")
                 except Exception as e:
-                    print(f"[LORA] Could not set scale: {e}")
+                    print(f"[LORA] Could not fuse: {e}")
             
             # ═══════════════════════════════════════════════════════════
             # PIPELINE CONFIGURATION

models.py

@@ -1,6 +1,6 @@
 """
 Model loading and initialization for Pixagram AI Pixel Art Generator
-ALTERNATIVE VERSION - Solution 3: Manually set controlnet after pipeline loading
+CORRECTED VERSION - Following examplewithface.py exactly
 """
 import torch
 import time
@@ -16,9 +16,10 @@ from controlnet_aux import ZoeDetector
 from huggingface_hub import hf_hub_download, snapshot_download
 from compel import Compel, ReturnedEmbeddingsType
 
-# Use InstantID pipeline
+# Use InstantID pipeline - import draw_kps too like in working example
 from pipeline_stable_diffusion_xl_instantid_img2img import (
-    StableDiffusionXLInstantIDImg2ImgPipeline
+    StableDiffusionXLInstantIDImg2ImgPipeline,
+    draw_kps
 )
 
 from config import (
@@ -175,39 +176,28 @@ def load_controlnets():
 def load_sdxl_pipeline(controlnets):
     """
     Load SDXL pipeline with InstantID support.
-    Solution 3: Load pipeline without controlnet, then set it manually.
+    controlnets MUST be a list: [identitynet, depthnet]
+    Following the working example from examplewithface.py
     """
     print("Loading SDXL checkpoint with InstantID pipeline...")
     try:
         model_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['checkpoint'])
         
-        # Load pipeline WITHOUT controlnet first (Solution 3)
-        print("  Loading pipeline from single file (without controlnet)...")
+        # Wrap multiple controlnets in MultiControlNetModel (CRITICAL!)
+        if isinstance(controlnets, list) and len(controlnets) > 1:
+            print(f"  Wrapping {len(controlnets)} ControlNets in MultiControlNetModel...")
+            multi_controlnet = MultiControlNetModel(controlnets)
+        else:
+            multi_controlnet = controlnets[0] if isinstance(controlnets, list) else controlnets
+        
+        # Use InstantID-enabled pipeline
         pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
             model_path,
-            controlnet=None,  # Load without controlnet initially
+            controlnet=multi_controlnet,
             torch_dtype=dtype,
             use_safetensors=True
         ).to(device)
         
-        # Now manually set the controlnet
-        if isinstance(controlnets, list) and len(controlnets) > 1:
-            print(f"  Manually setting MultiControlNetModel with {len(controlnets)} nets...")
-            pipe.controlnet = MultiControlNetModel(controlnets).to(device, dtype=dtype)
-            print(f"  [VERIFY] pipe.controlnet type: {type(pipe.controlnet)}")
-            print(f"  [VERIFY] pipe.controlnet has {len(pipe.controlnet.nets)} nets")
-        else:
-            print("  Setting single ControlNet...")
-            pipe.controlnet = (controlnets[0] if isinstance(controlnets, list) else controlnets).to(device, dtype=dtype)
-            print(f"  [VERIFY] pipe.controlnet type: {type(pipe.controlnet)}")
-        
-        # Verify it's the right type
-        from diffusers.models import ControlNetModel
-        if isinstance(pipe.controlnet, (ControlNetModel, MultiControlNetModel)):
-            print("  ✅ ControlNet properly attached to pipeline")
-        else:
-            print(f"  ⚠️  WARNING: Unexpected controlnet type: {type(pipe.controlnet)}")
-        
         # Load IP-Adapter weights for InstantID
         print("Loading IP-Adapter for InstantID...")
         ip_adapter_path = download_model_with_retry(
@@ -229,36 +219,73 @@ def load_sdxl_pipeline(controlnets):
         print("  [WARNING] Falling back to standard SDXL pipeline (no InstantID)")
         from diffusers import StableDiffusionXLControlNetImg2ImgPipeline
         
-        # Same approach for fallback
+        # Wrap controlnets for fallback too
+        if isinstance(controlnets, list) and len(controlnets) > 1:
+            multi_controlnet = MultiControlNetModel(controlnets)
+        else:
+            multi_controlnet = controlnets[0] if isinstance(controlnets, list) else controlnets
+        
         pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
             "stabilityai/stable-diffusion-xl-base-1.0",
-            controlnet=None,
+            controlnet=multi_controlnet,
             torch_dtype=dtype,
             use_safetensors=True
         ).to(device)
-        
-        # Manually set controlnet for fallback too
-        if isinstance(controlnets, list) and len(controlnets) > 1:
-            pipe.controlnet = MultiControlNetModel(controlnets).to(device, dtype=dtype)
-        else:
-            pipe.controlnet = (controlnets[0] if isinstance(controlnets, list) else controlnets).to(device, dtype=dtype)
-        
         return pipe, False
 
 
+# Global variable to track LORA fusion state (like in working example)
+last_lora_fused = False
+
+
 def load_lora(pipe):
-    """Load LORA from HuggingFace Hub."""
+    """
+    Load LORA from HuggingFace Hub.
+    Following working example: load_lora_weights then fuse_lora
+    """
     print("Loading LORA (retroart) from HuggingFace Hub...")
+    global last_lora_fused
+    
     try:
         lora_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['lora'])
-        pipe.load_lora_weights(lora_path, adapter_name="retroart")
-        print(f"  [OK] LORA loaded successfully")
+        
+        # Load the LORA weights (don't fuse yet - will fuse with scale before generation)
+        pipe.load_lora_weights(lora_path)
+        
+        print(f"  [OK] LORA loaded successfully (will fuse with scale before generation)")
+        last_lora_fused = False  # Track that we need to fuse
         return True
+        
     except Exception as e:
         print(f"  [WARNING] Could not load LORA: {e}")
         return False
 
 
+def fuse_lora_with_scale(pipe, lora_scale):
+    """
+    Fuse LORA with scale before generation (like in working example).
+    This should be called before each generation if scale changed.
+    """
+    global last_lora_fused
+    
+    try:
+        # Unfuse if already fused
+        if last_lora_fused:
+            print(f"  [LORA] Unfusing previous LORA...")
+            pipe.unfuse_lora()
+        
+        # Fuse with new scale
+        print(f"  [LORA] Fusing with scale: {lora_scale}")
+        pipe.fuse_lora(lora_scale)
+        last_lora_fused = True
+        
+        return True
+        
+    except Exception as e:
+        print(f"  [WARNING] Could not fuse LORA: {e}")
+        return False
+
+
 def setup_compel(pipe):
     """Setup Compel for better SDXL prompt handling."""
     print("Setting up Compel for enhanced prompt processing...")
@@ -353,4 +380,7 @@ def set_clip_skip(pipe):
         print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-print("[OK] Model loading functions ready (Solution 3: Manual controlnet setting)")
+# Export draw_kps for use in generator
+__all__ = ['draw_kps', 'fuse_lora_with_scale', 'last_lora_fused']
+
+print("[OK] Model loading functions ready (following examplewithface.py)")