Spaces:
Runtime error
Runtime error
| """ | |
| 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") |