Update app.py

app.py

@@ -3,55 +3,50 @@ import torch
 import os
 from diffusers import StableDiffusion3Pipeline
 from safetensors.torch import load_file
-from spaces import GPU  # Import GPU if in HF Space, otherwise remove this line
+from spaces import GPU  # Remove if not in HF Space
 
-# Access HF_TOKEN from environment variables
-hf_token = os.getenv("HF_TOKEN")
+# ... (HF_TOKEN, model_id - same as before)
 
-# Specify the pre-trained model ID
-model_id = "stabilityai/stable-diffusion-3.5-large"
-
-# Initialize pipeline *outside* the function (but set to None initially)
-pipeline = None
-
-
-# Function to load the Stable Diffusion pipeline (called only ONCE)
-def load_pipeline():
-    global pipeline
-    try:
-        pipeline = StableDiffusion3Pipeline.from_pretrained(
-            model_id,
-            use_auth_token=hf_token,
-            torch_dtype=torch.float16,
-            cache_dir="./model_cache"
-        )
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return f"Error loading model: {e}"
+pipeline = None  # Global pipeline variable
 
+# Load Stable Diffusion and LoRA *immediately* (before Gradio)
+try:
+    pipeline = StableDiffusion3Pipeline.from_pretrained(
+        model_id,
+        use_auth_token=hf_token,
+        torch_dtype=torch.float16,
+        cache_dir="./model_cache"
+    )
     pipeline.enable_model_cpu_offload()
     pipeline.enable_attention_slicing()
-    return "Model loaded successfully"
-
-@GPU(duration=65)  # Use GPU decorator (ONLY if in HF Space)
-def generate_image(prompt):
-    global pipeline
-    if pipeline is None:
-        return "Model not loaded. Please wait."
 
     lora_filename = "lora_trained_model.safetensors"
     lora_path = os.path.join("./", lora_filename)
     print(f"Loading LoRA from: {lora_path}")
 
-    try:
-        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)
-        else:
-            return f"Error: LoRA file not found at {lora_path}"
-    except Exception as e:
-        return f"Error loading LoRA: {e}"
+    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
+    else:
+        print(f"Error: LoRA file not found at {lora_path}")
+        exit() # Exit 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
+
+
+# Function for image generation (now much simpler)
+@GPU(duration=65)  # Use GPU decorator (ONLY if in HF Space)
+def generate_image(prompt):
+    global pipeline
+    if pipeline is None:  # This should never happen now
+        return "Error: Stable Diffusion model not loaded!"
 
     try:
         image = pipeline(prompt).images[0]
@@ -60,13 +55,11 @@ def generate_image(prompt):
         return f"Error generating image: {e}"
 
 
+# Create the Gradio interface (no "Load Model" button needed)
 with gr.Blocks() as demo:
     prompt_input = gr.Textbox(label="Prompt")
     image_output = gr.Image(label="Generated Image")
     generate_button = gr.Button("Generate")
-    load_model_button = gr.Button("Load Model")
-
-    load_model_button.click(fn=load_pipeline, outputs=load_model_button)
 
     generate_button.click(
         fn=generate_image,