Update app.py

app.py

@@ -1,65 +1,59 @@
-import torch
-import spaces
-from diffusers import StableDiffusion3Pipeline
-from huggingface_hub import login
 import os
 import gradio as gr
+import torch
+from diffusers import StableDiffusion3Pipeline
+import spaces
+import random
 
-# Retrieve the token from the environment variable
-token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
-if token:
-    login(token=token)  # Log in with the retrieved token
-else:
-    raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")
+# 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)
 
-# Load the Stable Diffusion 3.5 model with lower precision (float16) if GPU is available
-model_id = "stabilityai/stable-diffusion-3.5-large"
-pipe = StableDiffusion3Pipeline.from_pretrained(model_id)
+    image = pipeline(prompt, generator=generator).images[0]
+    return image
 
-# Check if GPU is available, then move the model to the appropriate device
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-pipe.to(device)
+# Device selection
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+# Load the Hugging Face token securely
+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
+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
+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)
+    pipeline.load_lora_weights(lora_state_dict)
+    print("✅ LoRA weights loaded successfully!")
+else:
+    print("⚠️ LoRA file not found! Running base model.")
 
-# Define the path to the LoRA model
-lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🖼️ LoRA Fine-Tuned SD 3.5 Image Generator")
 
-# Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
-def load_lora_model(pipe, lora_model_path):
-    # When loading the LoRA weights
-    lora_weights = torch.load(lora_model_path, map_location=device, weights_only=True)
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter Prompt", value="A woman in her 20s with expressive black eyes, graceful face, elegant body, standing on the beach at sunset. Photorealistic, highly detailed.")
+        seed_input = gr.Number(label="Seed (optional)", value=None)
     
-    # Check if the transformer folder has the necessary attributes
-    print(dir(pipe.transformer))  # List available attributes of the transformer (formerly 'unet')
-
-    # Apply weights to the transformer submodule
-    try:
-        for name, param in pipe.transformer.named_parameters():  # Accessing transformer parameters
-            if name in lora_weights:
-                param.data += lora_weights[name]
-    except AttributeError:
-        print("The model doesn't have 'transformer' attributes. Please check the model structure.")
-        # Add alternative handling or exit
+    generate_btn = gr.Button("Generate Image")
+    output_image = gr.Image(label="Generated Image")
 
-    return pipe
-
-# Load and apply the LoRA model weights
-pipe = load_lora_model(pipe, lora_model_path)
-
-# Use the @spaces.gpu decorator to ensure compatibility with GPU or CPU as needed
-@spaces.GPU(duration=65)  # This ensures GPU is allocated for 65 seconds
-def generate(prompt, seed=None):
-    generator = torch.manual_seed(seed) if seed is not None else None
-    # Generate the image using the prompt
-    image = pipe(prompt, height=512, width=512, generator=generator).images[0]
-    return image
+    generate_btn.click(generate_image, inputs=[prompt_input, seed_input], outputs=output_image)
 
-# Gradio interface
-iface = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Textbox(label="Enter your prompt"),  # For the prompt
-        gr.Number(label="Enter a seed (optional)", value=None),  # For the seed
-    ],
-    outputs="image"
-)
-iface.launch()
+# Launch Gradio App
+demo.launch()