Update app.py

app.py

@@ -1,49 +1,55 @@
 # app.py
-from diffusers import StableDiffusionPipeline, AutoPipelineForText2Image
+from diffusers import StableDiffusionPipeline
 import torch
 import gradio as gr
+from huggingface_hub import hf_hub_download # Import to download the specific file
 import os
 
-# Define the model ID for the specific checkpoint
-# This is the full ID of the repository where the safetensors file lives
-MODEL_ID = "uhralk/Indigo_Furry_mix"
-# This is the exact filename of the model checkpoint you want to load
-# Make sure this matches the filename on the Hugging Face Hub EXACTLY.
+# Define the base model and the specific checkpoint file
+BASE_MODEL = "runwayml/stable-diffusion-v1-5"
+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 or if fp16 causes issues
 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)
+
 # Load the model
 try:
-    # Attempt to load the model directly as a full pipeline from the checkpoint
-    # This assumes the checkpoint is a full model merge, not just LoRA weights
-    pipe = AutoPipelineForText2Image.from_pretrained(
-        MODEL_ID,
+    # 1. Load the base Stable Diffusion 1.5 pipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
+        BASE_MODEL,
         torch_dtype=dtype,
-        use_safetensors=True,
-        # Specify the exact filename within the repository
-        # This tells diffusers to look for this specific file as the main model weights
-        # without needing an external base model or explicit LoRA loading.
-        model_file=CHECKPOINT_FILENAME
+        use_safetensors=True
     )
     pipe.to(device)
-    print(f"Model loaded successfully on {device} using direct checkpoint loading.")
+    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.")
 
 except Exception as e:
-    print(f"Error loading model directly on {device}: {e}")
-    print("Attempting to load on CPU with float32 as a fallback (may be very slow)...")
+    print(f"Error loading model: {e}")
+    # Fallback to CPU if GPU loading fails or if no GPU, and retry loading
     try:
-        pipe = AutoPipelineForText2Image.from_pretrained(
-            MODEL_ID,
+        print("Attempting to load on CPU with float32 as a fallback...")
+        pipe = StableDiffusionPipeline.from_pretrained(
+            BASE_MODEL,
             torch_dtype=torch.float32,
-            use_safetensors=True,
-            model_file=CHECKPOINT_FILENAME
+            use_safetensors=True
         )
         pipe.to("cpu")
-        print("Model forced loaded on 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.")
+
     except Exception as e_cpu:
         print(f"Critical error: Failed to load model even on CPU: {e_cpu}")
         exit() # Exit if model cannot be loaded at all
@@ -54,10 +60,13 @@ def generate_image(prompt, negative_prompt, num_inference_steps, guidance_scale,
     if not prompt:
         return None, "Please enter a prompt."
 
-    # Set up random seed if not -1
     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,
@@ -89,8 +98,7 @@ iface = gr.Interface(
         ["An indigo furry warrior, wielding a glowing sword, intricate armor, epic fantasy art", "blurry, low quality"],
         ["A cute indigo furry wizard casting a spell, magical effects, cartoon style", "disfigured, ugly"],
     ]
-) # <--- THIS CLOSING PARENTHESIS IS CRUCIAL AND OFTEN MISSED!
+)
 
-# Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    iface.