Upload 2 files

app.py

@@ -10,7 +10,6 @@ import torch
 import gradio as gr
 import spaces
 from tqdm.auto import tqdm
-import numpy as np
 from PIL import Image
 from utils import (
     load_models, clear_gpu_memory, set_timesteps, latents_to_pil, 
@@ -18,7 +17,6 @@ from utils import (
 )
 from diffusers import StableDiffusionPipeline
 
-
 # Set device
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 if device == "mps":
@@ -31,12 +29,11 @@ def load_model():
         "runwayml/stable-diffusion-v1-5",
         torch_dtype=torch.float16,
         safety_checker=None
-    )
+    ).to(device)
 
 @spaces.GPU
 def get_pipeline():
-    pipe = load_model()
-    return pipe.to("cuda")
+    return load_model()
 
 # Load concept library
 concept_embeds, concept_tokens = load_concept_library(get_pipeline())
@@ -50,6 +47,7 @@ art_concepts = {
     "comic_book": "comic book style, ink outlines, cel shading"
 }
 
+@spaces.GPU
 def generate_latents(prompt, seed, num_inference_steps, guidance_scale, 
                      vignette_loss_scale, concept_style=None, concept_strength=0.5,
                      height=512, width=512):
@@ -267,7 +265,6 @@ def create_demo():
                     guidance_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.1, label="Guidance Scale", value=7.5)
                     vignette_loss_scale = gr.Slider(minimum=0.0, maximum=100.0, step=1.0, label="Vignette Loss Scale", value=0.0)
                     
-                    # Combine SD concept library tokens and art concept descriptions
                     all_styles = ["none"] + concept_tokens + list(art_concepts.keys())
                     concept_style = gr.Dropdown(choices=all_styles, label="Style Concept", value="none")
                     concept_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label="Concept Strength", value=0.5)
@@ -294,7 +291,7 @@ def create_demo():
         
         # Set up event handlers
         generate_btn.click(
-            generate_image,
+            generate_latents,
             inputs=[prompt, seed, num_inference_steps, guidance_scale, 
                     vignette_loss_scale, concept_style, concept_strength],
             outputs=output_image

utils.py

@@ -67,12 +67,11 @@ def clear_gpu_memory():
     """Clear GPU memory cache"""
     torch.cuda.empty_cache()
     gc.collect()
-    torch.cuda.empty_cache()
 
 def set_timesteps(scheduler, num_inference_steps):
     """Set timesteps for the scheduler with MPS compatibility fix"""
     scheduler.set_timesteps(num_inference_steps)
-    scheduler.timesteps = scheduler.timesteps.to(torch.float32)  # minor fix to ensure MPS compatibility
+    scheduler.timesteps = scheduler.timesteps.to(torch.float32)
 
 def pil_to_latent(input_im, vae, device):
     """