Update for ZeroGPU

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces
 import torch
 from torchvision import transforms
 from SDXL.diff_pipe import StableDiffusionXLDiffImg2ImgPipeline
@@ -9,7 +10,7 @@ device = "cuda"
 
 base = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
-)
+).to(device)
 
 refiner = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-refiner-1.0",
@@ -18,7 +19,7 @@ refiner = StableDiffusionXLDiffImg2ImgPipeline.from_pretrained(
     torch_dtype=torch.float16,
     use_safetensors=True,
     variant="fp16",
-)
+).to(device)
 
 base.scheduler = DPMSolverMultistepScheduler.from_config(base.scheduler.config)
 refiner.scheduler = DPMSolverMultistepScheduler.from_config(base.scheduler.config)
@@ -42,24 +43,21 @@ def preprocess_map(map):
     return map
 
 
+@spaces.GPU
 def inference(image, map, gs, prompt, negative_prompt):
     validate_inputs(image, map)
     image = preprocess_image(image)
     map = preprocess_map(map)
-    base_cuda = base.to(device)
-    edited_images = base_cuda(prompt=prompt, original_image=image, image=image, strength=1, guidance_scale=gs,
+    edited_images = base(prompt=prompt, original_image=image, image=image, strength=1, guidance_scale=gs,
                          num_images_per_prompt=1,
                          negative_prompt=negative_prompt,
                          map=map,
                          num_inference_steps=NUM_INFERENCE_STEPS, denoising_end=0.8, output_type="latent").images
-    base_cuda=None
-    refiner_cuda = refiner.to(device)
-    edited_images = refiner_cuda(prompt=prompt, original_image=image, image=edited_images, strength=1, guidance_scale=7.5,
+    edited_images = refiner(prompt=prompt, original_image=image, image=edited_images, strength=1, guidance_scale=7.5,
                             num_images_per_prompt=1,
                             negative_prompt=negative_prompt,
                             map=map,
                             num_inference_steps=NUM_INFERENCE_STEPS, denoising_start=0.8).images[0]
-    refiner_cuda=None
     return edited_images