""" Model loading and initialization for Pixagram AI Pixel Art Generator UPDATED VERSION with proper InstantID pipeline support """ import torch import time from diffusers import ( ControlNetModel, AutoencoderKL, LCMScheduler ) from insightface.app import FaceAnalysis from controlnet_aux import ZoeDetector from huggingface_hub import hf_hub_download from compel import Compel, ReturnedEmbeddingsType # Use InstantID pipeline from pipeline_stable_diffusion_xl_instantid_img2img import ( StableDiffusionXLInstantIDImg2ImgPipeline ) from config import ( device, dtype, MODEL_REPO, MODEL_FILES, HUGGINGFACE_TOKEN, FACE_DETECTION_CONFIG, CLIP_SKIP, DOWNLOAD_CONFIG ) def download_model_with_retry(repo_id, filename, max_retries=None): """Download model with retry logic and proper token handling.""" 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})...") 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] Download 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") raise return None def load_face_analysis(): """Load face analysis model on CPU to save GPU memory.""" print("Loading face analysis model on CPU...") try: # Force CPU execution for face analysis to save GPU memory face_app = FaceAnalysis( name=FACE_DETECTION_CONFIG['model_name'], root='./models/insightface', providers=['CPUExecutionProvider'] # CPU only for face detection ) face_app.prepare( ctx_id=-1, # -1 for CPU det_size=FACE_DETECTION_CONFIG['det_size'] ) print(" [OK] Face analysis model loaded on CPU (GPU memory saved)") return face_app, True except Exception as e: print(f" [WARNING] Face detection not available: {e}") return None, False def load_depth_detector(): """Load Zoe Depth detector with optimized memory management.""" print("Loading Zoe Depth detector...") try: zoe_depth = ZoeDetector.from_pretrained("lllyasviel/Annotators") # Start on CPU to save memory during initialization zoe_depth = zoe_depth.to("cpu") print(" [OK] Zoe Depth loaded (on CPU, will move to GPU when needed)") return zoe_depth, True except Exception as e: print(f" [WARNING] Zoe Depth not available: {e}") return None, False def load_controlnets(): """ Load ControlNets for InstantID pipeline. Returns both ControlNets (InstantID first, then Depth). """ print("Loading InstantID ControlNet...") controlnet_instantid = ControlNetModel.from_pretrained( "InstantX/InstantID", subfolder="ControlNetModel", torch_dtype=dtype ).to(device) print(" [OK] InstantID ControlNet loaded") print("Loading Zoe Depth ControlNet...") controlnet_depth = ControlNetModel.from_pretrained( "diffusers/controlnet-zoe-depth-sdxl-1.0", torch_dtype=dtype ).to(device) print(" [OK] Zoe Depth ControlNet loaded") return controlnet_instantid, controlnet_depth def load_sdxl_pipeline(controlnets): """ Load SDXL pipeline with InstantID support. controlnets MUST be a list: [identitynet, depthnet] """ print("Loading SDXL checkpoint with InstantID pipeline...") try: model_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['checkpoint']) # Use InstantID-enabled pipeline pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file( model_path, controlnet=controlnets, torch_dtype=dtype, use_safetensors=True ).to(device) # Load IP-Adapter weights for InstantID 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) pipe.set_ip_adapter_scale(0.8) # Default scale print(" [OK] InstantID pipeline loaded successfully") return pipe, True except Exception as e: print(f" [ERROR] Could not load InstantID pipeline: {e}") import traceback traceback.print_exc() # Fallback to standard pipeline print(" Falling back to standard SDXL pipeline (no InstantID)") from diffusers import StableDiffusionXLControlNetImg2ImgPipeline pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", controlnet=controlnets, torch_dtype=dtype, use_safetensors=True ).to(device) return pipe, False def load_lora(pipe): """Load LORA from HuggingFace Hub.""" print("Loading LORA (retroart) from HuggingFace Hub...") 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") return True except Exception as e: print(f" [WARNING] Could not load LORA: {e}") return False def setup_compel(pipe): """Setup Compel for better SDXL prompt handling.""" print("Setting up Compel for enhanced prompt processing...") try: compel = Compel( tokenizer=[pipe.tokenizer, pipe.tokenizer_2], text_encoder=[pipe.text_encoder, pipe.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True] ) print(" [OK] Compel loaded successfully") return compel, True except Exception as e: print(f" [WARNING] Compel not available: {e}") return None, False def setup_scheduler(pipe): """Setup LCM scheduler.""" print("Setting up LCM scheduler...") pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) print(" [OK] LCM scheduler configured") def optimize_pipeline(pipe): """Apply optimizations to pipeline.""" 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}") def load_caption_model(): """ Load caption model with proper error handling. Tries multiple models in order of quality. Models start on CPU and move to GPU only when needed. """ print("Loading caption model...") # Try GIT-Large first try: from transformers import AutoProcessor, AutoModelForCausalLM print(" Attempting GIT-Large (recommended)...") caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco") caption_model = AutoModelForCausalLM.from_pretrained( "microsoft/git-large-coco", torch_dtype=dtype # Use dtype from config ).to("cpu") # Start on CPU to save GPU memory print(" [OK] GIT-Large model loaded (on CPU, will move to GPU when needed)") return caption_processor, caption_model, True, 'git' except Exception as e1: print(f" [INFO] GIT-Large not available: {e1}") # Try BLIP base as fallback try: from transformers import BlipProcessor, BlipForConditionalGeneration print(" Attempting BLIP base (fallback)...") caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") caption_model = BlipForConditionalGeneration.from_pretrained( "Salesforce/blip-image-captioning-base", torch_dtype=dtype # Use dtype from config ).to("cpu") # Start on CPU to save GPU memory print(" [OK] BLIP base model loaded (on CPU, will move to GPU when needed)") return caption_processor, caption_model, True, 'blip' except Exception as e2: print(f" [WARNING] Caption models not available: {e2}") return None, None, False, 'none' def set_clip_skip(pipe): """Set CLIP skip value.""" if hasattr(pipe, 'text_encoder'): print(f" [OK] CLIP skip set to {CLIP_SKIP}") print("[OK] Model loading functions ready")