Update app.py

app.py

@@ -17,7 +17,7 @@ pipeline = None
 
 # Function to load the Stable Diffusion pipeline (called only ONCE)
 def load_pipeline():
-    global pipeline  # Use the global keyword to modify the global variable
+    global pipeline
     try:
         pipeline = StableDiffusion3Pipeline.from_pretrained(
             model_id,
@@ -27,31 +27,31 @@ def load_pipeline():
         )
     except Exception as e:
         print(f"Error loading model: {e}")
-        return f"Error loading model: {e}"  # Return error message
+        return f"Error loading model: {e}"
 
     pipeline.enable_model_cpu_offload()
     pipeline.enable_attention_slicing()
-    return "Model loaded successfully" # Return success message
+    return "Model loaded successfully"
 
-# 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:  # Check if pipeline is loaded
-        return "Model not loaded. Please wait." # Return message if not loaded
+    if pipeline is None:
+        return "Model not loaded. Please wait."
 
-    # Load and apply LoRA (file is already in the Space)
-    lora_filename = "lora_trained_model.safetensors"  # Name of your LoRA file
-    lora_path = os.path.join("./", lora_filename)  # Construct the path
+    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):  # check if the file exists
+        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)
-        except Exception as e:
-            return f"Error loading LoRA: {e}"
+        else:
+            return f"Error: LoRA file not found at {lora_path}"
+    except Exception as e:
+        return f"Error loading LoRA: {e}"
 
     try:
         image = pipeline(prompt).images[0]
@@ -60,14 +60,13 @@ def generate_image(prompt):
         return f"Error generating image: {e}"
 
 
-# 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 = gr.Button("Load Model")
 
-    load_model_button.click(fn=load_pipeline, outputs=load_model_button) # Call load_pipeline
+    load_model_button.click(fn=load_pipeline, outputs=load_model_button)
 
     generate_button.click(
         fn=generate_image,