Update app.py

app.py

@@ -4,18 +4,6 @@ import torch
 from diffusers import StableDiffusion3Pipeline
 import spaces
 import random
-from peft import PeftModel, get_peft_model
-
-# Ensure GPU allocation in Hugging Face Spaces
-@spaces.GPU(duration=65)
-def generate_image(prompt: str, seed: int = None):
-    """Generates an image using Stable Diffusion 3.5 with LoRA fine-tuning."""
-    if seed is None:
-        seed = random.randint(0, 100000)
-    generator = torch.manual_seed(seed)
-
-    image = pipeline(prompt, generator=generator).images[0]
-    return image
 
 # Device selection
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -27,22 +15,36 @@ token = os.getenv("HF_TOKEN")
 # Model ID for SD 3.5 Large
 model_repo_id = "stabilityai/stable-diffusion-3.5-large"
 
-# Load Stable Diffusion pipeline
+# Load Stable Diffusion pipeline once at the start
 pipeline = StableDiffusion3Pipeline.from_pretrained(
     model_repo_id,
     torch_dtype=torch_dtype,
     use_safetensors=True,  # Use safetensors format if supported
 ).to(device)
 
-# Load the LoRA trained weights
+# Load the LoRA trained weights once at the start
 lora_path = "lora_trained_model.pt"  # Ensure this file is uploaded in the Space
 if os.path.exists(lora_path):
     lora_state_dict = torch.load(lora_path, map_location=device, weights_only=True)
-    pipeline = PeftModel.from_pretrained(pipeline, lora_path)
+    pipeline.load_lora_weights(lora_state_dict)  # Load LoRA weights into the pipeline
     print("✅ LoRA weights loaded successfully!")
 else:
     print("⚠️ LoRA file not found! Running base model.")
 
+# Ensure GPU allocation in Hugging Face Spaces
+@spaces.GPU(duration=65)
+def generate_image(prompt: str, seed: int = None):
+    """Generates an image using Stable Diffusion 3.5 with LoRA fine-tuning."""
+    if seed is None:
+        seed = random.randint(0, 100000)
+    
+    # Create a generator with the seed
+    generator = torch.manual_seed(seed)
+
+    # Generate the image using the pipeline
+    image = pipeline(prompt, generator=generator).images[0]
+    return image
+
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("# 🖼️ LoRA Fine-Tuned SD 3.5 Image Generator")