Update app.py

app.py

@@ -786,66 +786,38 @@ def load_illustrious_xl(
     filename: str = "illustriousXL_v01.safetensors",
     device: str = "cuda"
 ) -> Tuple[UNet2DConditionModel, AutoencoderKL, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer, CLIPTokenizer]:
-    """Load Illustrious XL from single safetensors file."""
+    """Load Illustrious XL from single safetensors file using diffusers' single-file loader."""
+    from diffusers import StableDiffusionXLPipeline
     
-    print(f"📥 Downloading Illustrious XL: {repo_id}/{filename}")
+    print(f"📥 Loading Illustrious XL: {repo_id}/{filename}")
     
+    # Download the checkpoint
     checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
     print(f"✓ Downloaded: {checkpoint_path}")
     
-    print("📦 Loading safetensors...")
-    state_dict = load_safetensors(checkpoint_path)
-    
-    # Extract components
-    components = extract_comfyui_components(state_dict)
-    
-    # Load UNet from SDXL base config, then load weights
-    print("🏗️ Initializing SDXL UNet...")
-    unet = UNet2DConditionModel.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        subfolder="unet",
-        torch_dtype=torch.float16
+    # Use diffusers' built-in single-file loader which handles key remapping
+    print("📦 Loading with StableDiffusionXLPipeline.from_single_file()...")
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        checkpoint_path,
+        torch_dtype=torch.float16,
+        use_safetensors=True,
     )
     
-    if components["unet"]:
-        missing, unexpected = unet.load_state_dict(components["unet"], strict=False)
-        print(f"  UNet: {len(missing)} missing, {len(unexpected)} unexpected keys")
-    
-    # Load VAE
-    print("🏗️ Initializing SDXL VAE...")
-    vae = AutoencoderKL.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        subfolder="vae",
-        torch_dtype=torch.float16
-    )
-    
-    if components["vae"]:
-        missing, unexpected = vae.load_state_dict(components["vae"], strict=False)
-        print(f"  VAE: {len(missing)} missing, {len(unexpected)} unexpected keys")
-    
-    # Load CLIP-L
-    print("🏗️ Loading CLIP-L...")
-    text_encoder = CLIPTextModel.from_pretrained(
-        "openai/clip-vit-large-patch14",
-        torch_dtype=torch.float16
-    )
-    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-    
-    # Load CLIP-G
-    print("🏗️ Loading CLIP-G...")
-    text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-        "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
-        torch_dtype=torch.float16
-    )
-    tokenizer_2 = CLIPTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
+    # Extract components
+    unet = pipe.unet.to(device)
+    vae = pipe.vae.to(device)
+    text_encoder = pipe.text_encoder.to(device)
+    text_encoder_2 = pipe.text_encoder_2.to(device)
+    tokenizer = pipe.tokenizer
+    tokenizer_2 = pipe.tokenizer_2
     
-    # Move to device
-    unet = unet.to(device)
-    vae = vae.to(device)
-    text_encoder = text_encoder.to(device)
-    text_encoder_2 = text_encoder_2.to(device)
+    # Clean up the pipeline to free memory
+    del pipe
+    torch.cuda.empty_cache()
     
     print("✅ Illustrious XL loaded!")
+    print(f"   UNet params: {sum(p.numel() for p in unet.parameters()):,}")
+    print(f"   VAE params: {sum(p.numel() for p in vae.parameters()):,}")
     
     return unet, vae, text_encoder, text_encoder_2, tokenizer, tokenizer_2