Update app.py

app.py

@@ -1,19 +1,9 @@
-import torch
-import zerogpu  # Import ZeroGPU
+import spaces
 from diffusers import StableDiffusion3Pipeline
 from huggingface_hub import login
 import os
 import gradio as gr
 
-# Automatically choose GPU if available, otherwise CPU
-device = zerogpu.select_device()  # ZeroGPU will automatically choose 'cuda' or 'cpu'
-
-# Check and print if the selected device is GPU or CPU
-if device == "cuda":
-    print(f"Using GPU: {torch.cuda.get_device_name()}")
-else:
-    print("Using CPU")
-
 # Retrieve the token from the environment variable
 token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
 if token:
@@ -23,12 +13,10 @@ else:
 
 # Load the Stable Diffusion 3.5 model with lower precision (float16) if GPU is available
 model_id = "stabilityai/stable-diffusion-3.5-large"
-if device == "cuda":
-    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 precision
-else:
-    pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # Default precision for CPU
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id)
 
-pipe.to(device)  # Ensuring the model is on the correct device (GPU or CPU)
+# Check if GPU is available, then move the model to the appropriate device
+pipe.to('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Define the path to the LoRA model
 lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
@@ -36,7 +24,7 @@ lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
 # Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
 def load_lora_model(pipe, lora_model_path):
     # Load the LoRA weights
-    lora_weights = torch.load(lora_model_path, map_location=device)  # Load LoRA model to the correct device
+    lora_weights = torch.load(lora_model_path, map_location=pipe.device)  # Load LoRA model to the correct device
     
     # Print available attributes of the model to check access to `unet` (optional)
     print(dir(pipe))  # This will list all attributes and methods of the `pipe` object
@@ -55,16 +43,17 @@ def load_lora_model(pipe, lora_model_path):
 # Load and apply the LoRA model weights
 pipe = load_lora_model(pipe, lora_model_path)
 
-# Function to generate an image from a text prompt
-def generate_image(prompt, seed=None):
+# Use the @space.gpu decorator to ensure compatibility with GPU or CPU as needed
+@spaces.gpu
+def generate(prompt, seed=None):
     generator = torch.manual_seed(seed) if seed is not None else None
-    # Reduce image size for less memory usage
-    image = pipe(prompt, height=512, width=512, generator=generator).images[0]  # Changed image size
+    # Generate the image using the prompt
+    image = pipe(prompt, height=512, width=512, generator=generator).images[0]
     return image
 
 # Gradio interface
 iface = gr.Interface(
-    fn=generate_image,
+    fn=generate,
     inputs=[
         gr.Textbox(label="Enter your prompt"),  # For the prompt
         gr.Number(label="Enter a seed (optional)", value=None),  # For the seed