Update app.py

app.py

@@ -9,14 +9,14 @@ base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_4step_unet.pth"
 
-# Model Initialization
-pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
-pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
-pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-
 # Function 
 @spaces.GPU
 def generate_image(prompt):
+    # Ensure model and scheduler are initialized in GPU-enabled function
+    pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+    pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, ckpt), map_location="cuda"))
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
     image = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
     return image
 
@@ -34,5 +34,4 @@ demo = gr.Interface(
     description=description
 )
 
-demo.queue(max_size=20)
-demo.launch()
+demo.launch()