Update app.py

app.py

@@ -2,56 +2,64 @@
 from diffusers import StableDiffusionPipeline
 import torch
 import gradio as gr
-from huggingface_hub import hf_hub_download # Import to download the specific file
+from huggingface_hub import hf_hub_download
 import os
 
-# Define the base model and the specific checkpoint file
-BASE_MODEL = "runwayml/stable-diffusion-v1-5"
+# Define the model ID for the specific checkpoint
+# We will use this to get the direct download URL for from_single_file
 MODEL_REPO_ID = "uhralk/Indigo_Furry_mix"
 CHECKPOINT_FILENAME = "indigo_Furrymix_v120_hybrid_fin_fp16.safetensors"
 
 # Determine the device (GPU or CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
+# Use float16 for GPU to save VRAM and speed up, float32 for CPU
+# Stable Diffusion models are often best run in fp16 on GPU
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Path to the downloaded checkpoint file
-# This will download the file to the Hugging Face cache directory
-checkpoint_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME)
+# Construct the full Hub URL for from_single_file
+# The format is "hf_hub_url/repo_id/filename"
+full_checkpoint_url = f"hf_hub:{MODEL_REPO_ID}/{CHECKPOINT_FILENAME}"
 
 # Load the model
 try:
-    # 1. Load the base Stable Diffusion 1.5 pipeline
-    pipe = StableDiffusionPipeline.from_pretrained(
-        BASE_MODEL,
+    print(f"Attempting to load model from single file: {full_checkpoint_url} on {device}")
+    # Load the pipeline directly from the single .safetensors file
+    pipe = StableDiffusionPipeline.from_single_file(
+        full_checkpoint_url,
         torch_dtype=dtype,
-        use_safetensors=True
+        use_safetensors=True,
+        # Ensure VAE, Text Encoder, and Scheduler are fetched from a compatible base if not included
+        # SD 1.5 components are compatible with most fine-tunes
+        vae=None, # Automatically inferred or loaded from base
+        text_encoder=None, # Automatically inferred or loaded from base
+        tokenizer=None, # Automatically inferred or loaded from base
+        scheduler=None, # Automatically inferred or loaded from base
+        # Specifying a compatible base ensures non-Unet components are loaded
+        from_transformers=False, # Important for diffusers checkpoints
+        load_safety_checker=True # Good practice
     )
     pipe.to(device)
-    print(f"Base model '{BASE_MODEL}' loaded successfully on {device}")
-
-    # 2. Load the custom checkpoint's state dict into the pipeline
-    # This is often the most compatible way for "merged" checkpoints
-    pipe.load_lora_weights(checkpoint_path, adapter_name="indigo_mix") # Use a dummy adapter name
-    print(f"Custom checkpoint '{CHECKPOINT_FILENAME}' loaded successfully.")
+    print(f"Model loaded successfully from single file on {device}")
 
 except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fallback to CPU if GPU loading fails or if no GPU, and retry loading
+    print(f"Error loading model directly from single file on {device}: {e}")
+    print("Attempting to load on CPU with float32 as a fallback (may be very slow)...")
     try:
-        print("Attempting to load on CPU with float32 as a fallback...")
-        pipe = StableDiffusionPipeline.from_pretrained(
-            BASE_MODEL,
-            torch_dtype=torch.float32,
-            use_safetensors=True
+        pipe = StableDiffusionPipeline.from_single_file(
+            full_checkpoint_url,
+            torch_dtype=torch.float32, # Force float32 for CPU
+            use_safetensors=True,
+            vae=None,
+            text_encoder=None,
+            tokenizer=None,
+            scheduler=None,
+            from_transformers=False,
+            load_safety_checker=True
         )
         pipe.to("cpu")
-        print(f"Base model '{BASE_MODEL}' loaded successfully on CPU")
-
-        pipe.load_lora_weights(checkpoint_path, adapter_name="indigo_mix") # Load on CPU as well
-        print(f"Custom checkpoint '{CHECKPOINT_FILENAME}' loaded successfully on CPU.")
-
+        print("Model forced loaded on CPU.")
     except Exception as e_cpu:
-        print(f"Critical error: Failed to load model even on CPU: {e_cpu}")
+        print(f"Critical error: Failed to load model even from single file on CPU: {e_cpu}")
         exit() # Exit if model cannot be loaded at all
 
 
@@ -63,10 +71,6 @@ def generate_image(prompt, negative_prompt, num_inference_steps, guidance_scale,
     generator = torch.Generator(device=device).manual_seed(seed) if seed != -1 else None
 
     try:
-        # If you loaded the LoRA/merged checkpoint, you might need to enable it for inference
-        # If it was loaded as a direct state dict, it's already "active"
-        # For pipe.load_lora_weights, if not using a specific adapter name for later enable/disable:
-        # you just call it and it applies immediately.
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -99,6 +103,6 @@ iface = gr.Interface(
         ["A cute indigo furry wizard casting a spell, magical effects, cartoon style", "disfigured, ugly"],
     ]
 )
-# Launch the Gradio app
+
 if __name__ == "__main__":
-    iface.launch() # <--- Make sure this line is exactly 'iface.launch()'
+    iface.launch()