Update app.py

app.py

@@ -1,42 +1,52 @@
 # app.py
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, AutoPipelineForText2Image
 import torch
 import gradio as gr
 import os
 
-# 1. Define the base model and the checkpoint model
-BASE_MODEL = "runwayml/stable-diffusion-v1-5"
-CHECKPOINT_MODEL = "uhralk/Indigo_Furry_mix"
-CHECKPOINT_FILENAME = "indigo_Furrymix_v120_hybrid_fin_fp16.safetensors" # Using the recommended filename
+# 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.
+CHECKPOINT_FILENAME = "indigo_Furrymix_v120_hybrid_fin_fp16.safetensors"
 
-# 2. Load the model
+# 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
 
+# Load the model
 try:
-    pipe = StableDiffusionPipeline.from_pretrained(
-        BASE_MODEL,
+    # 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,
         torch_dtype=dtype,
-        use_safetensors=True
+        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
     )
-    pipe.load_lora_weights(CHECKPOINT_MODEL, weight_name=CHECKPOINT_FILENAME, subfolder="")
     pipe.to(device)
-    print(f"Model loaded successfully on {device}")
+    print(f"Model loaded successfully on {device} using direct checkpoint loading.")
+
 except Exception as e:
-    print(f"Error loading model on {device}: {e}")
-    # Fallback to CPU if GPU loading fails or if no GPU
+    print(f"Error loading model directly on {device}: {e}")
+    print("Attempting to load on CPU with float32 as a fallback (may be very slow)...")
     try:
-        pipe = StableDiffusionPipeline.from_pretrained(
-            BASE_MODEL,
+        pipe = AutoPipelineForText2Image.from_pretrained(
+            MODEL_ID,
             torch_dtype=torch.float32,
-            use_safetensors=True
+            use_safetensors=True,
+            model_file=CHECKPOINT_FILENAME
         )
-        pipe.load_lora_weights(CHECKPOINT_MODEL, weight_name=CHECKPOINT_FILENAME, subfolder="")
         pipe.to("cpu")
-        print("Model forced loaded on CPU due to previous error or no GPU.")
+        print("Model forced loaded on CPU.")
     except Exception as e_cpu:
         print(f"Critical error: Failed to load model even on CPU: {e_cpu}")
-        exit()
+        exit() # Exit if model cannot be loaded at all
 
 
 # 3. Define the Gradio inference function
@@ -44,6 +54,7 @@ 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:
@@ -72,13 +83,4 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Generated Image"),
         gr.Textbox(label="Status/Error")
     ],
-    title="Indigo Furry Mix - Text-to-Image Generator",
-    description="Generate images of furry characters with the Indigo Furry Mix model.",
-    examples=[
-        ["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"],
-    ]
-)
-
-if __name__ == "__main__":
-    iface.launch()
+    title="Indigo Fur