Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import os
 from diffusers import StableDiffusion3Pipeline
 from safetensors.torch import load_file
-from spaces import GPU 
+from spaces import GPU  # Import GPU if in HF Space, otherwise remove this line
 
 # Access HF_TOKEN from environment variables
 hf_token = os.getenv("HF_TOKEN")
@@ -11,31 +11,34 @@ hf_token = os.getenv("HF_TOKEN")
 # Specify the pre-trained model ID
 model_id = "stabilityai/stable-diffusion-3.5-large"
 
-# Lazy pipeline initialization
+# Initialize pipeline *outside* the function (but set to None initially)
 pipeline = None
 
-# Function for image generation
-@gr.GPU(duration=65)
-def generate_image(prompt):  # Remove lora_file input
+
+# Function to load the Stable Diffusion pipeline (called only ONCE)
+def load_pipeline():
+    global pipeline  # Use the global keyword to modify the global variable
+    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}"  # Return error message
+
+    pipeline.enable_model_cpu_offload()
+    pipeline.enable_attention_slicing()
+    return "Model loaded successfully" # Return success message
+
+# Function for image generation (now decorated)
+@GPU(duration=65)  # Use GPU decorator (ONLY if in HF Space)
+def generate_image(prompt):
     global pipeline
-    if pipeline is None:
-        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 from cache: {e}")
-            pipeline = StableDiffusion3Pipeline.from_pretrained(
-                model_id,
-                use_auth_token=hf_token,
-                torch_dtype=torch.float16,
-                local_files_only=False
-            )
-        pipeline.enable_model_cpu_offload()
-        pipeline.enable_attention_slicing()
+    if pipeline is None:  # Check if pipeline is loaded
+        return "Model not loaded. Please wait." # Return message if not loaded
 
     # Load and apply LoRA (file is already in the Space)
     lora_filename = "lora_trained_model.safetensors"  # Name of your LoRA file
@@ -47,10 +50,8 @@ def generate_image(prompt):  # Remove lora_file input
             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}"
+        except Exception as e:
+            return f"Error loading LoRA: {e}"
 
     try:
         image = pipeline(prompt).images[0]
@@ -59,15 +60,18 @@ def generate_image(prompt):  # Remove lora_file input
         return f"Error generating image: {e}"
 
 
-# Create the Gradio interface (remove lora_upload)
+# Create the Gradio interface
 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") # Button to load model
+
+    load_model_button.click(fn=load_pipeline, outputs=load_model_button) # Call load_pipeline
 
     generate_button.click(
         fn=generate_image,
-        inputs=prompt_input,  # Only prompt input now
+        inputs=prompt_input,
         outputs=image_output,
     )