Update app.py

app.py

@@ -3,59 +3,56 @@ import torch
 import os
 from diffusers import StableDiffusion3Pipeline
 from safetensors.torch import load_file
-from spaces import GPU  # Remove if not in HF Space
+from spaces import GPU  # Remove this line if NOT in a HF Space
 
-# ... (HF_TOKEN, model_id - same as before)
+# 1. Define model ID and HF_TOKEN (at the VERY beginning)
+model_id = "stabilityai/stable-diffusion-3.5-large"  # Correct model ID for SD 3.5 Large
+hf_token = os.getenv("HF_TOKEN")  # For private models (set in HF Space settings)
 
-pipeline = None  # Global pipeline variable
+# 2. Initialize pipeline (to None initially)
+pipeline = None
 
-# Load Stable Diffusion and LoRA *immediately* (before Gradio)
+# 3. Load Stable Diffusion and LoRA (before Gradio)
 try:
     pipeline = StableDiffusion3Pipeline.from_pretrained(
         model_id,
         use_auth_token=hf_token,
-        torch_dtype=torch.float16,
-        cache_dir="./model_cache"
+        torch_dtype=torch.float16,  # Use float16 for memory efficiency
+        cache_dir="./model_cache"  # For caching
     )
-    pipeline.enable_model_cpu_offload()
-    pipeline.enable_attention_slicing()
 
-    lora_filename = "lora_trained_model.safetensors"
+    lora_filename = "lora_trained_model.safetensors"  # EXACT filename of your LoRA
     lora_path = os.path.join("./", lora_filename)
-    print(f"Loading LoRA from: {lora_path}")
 
     if os.path.exists(lora_path):
         lora_weights = load_file(lora_path)
         text_encoder = pipeline.text_encoder
         text_encoder.load_state_dict(lora_weights, strict=False)
-        print("LoRA loaded successfully!") # Confirmation message
+        print(f"LoRA loaded successfully from: {lora_path}")
     else:
-        print(f"Error: LoRA file not found at {lora_path}")
-        exit() # Exit if LoRA is not found
-
+        print(f"Error: LoRA file not found at: {lora_path}")
+        exit()  # Stop if LoRA is not found
 
     print("Stable Diffusion model loaded successfully!")
 
 except Exception as e:
-    print(f"Error loading Stable Diffusion or LoRA: {e}")
-    exit()  # Exit if there's an error
-
+    print(f"Error loading model or LoRA: {e}")
+    exit()  # Stop if model loading fails
 
-# Function for image generation (now much simpler)
-@GPU(duration=65)  # Use GPU decorator (ONLY if in HF Space)
+# 4. Image generation function (now decorated)
+@GPU(duration=65)  # ONLY if in a HF Space, remove if not
 def generate_image(prompt):
     global pipeline
-    if pipeline is None:  # This should never happen now
-        return "Error: Stable Diffusion model not loaded!"
+    if pipeline is None:  # Should not happen, but good to check
+        return "Error: Model not loaded!"
 
     try:
-        image = pipeline(prompt).images[0]
+        image = pipeline(prompt).images[0]  # Access the first image from the list
         return image
     except Exception as e:
         return f"Error generating image: {e}"
 
-
-# Create the Gradio interface (no "Load Model" button needed)
+# 5. Gradio interface
 with gr.Blocks() as demo:
     prompt_input = gr.Textbox(label="Prompt")
     image_output = gr.Image(label="Generated Image")