Update app.py

app.py

@@ -7,36 +7,42 @@ import spaces
 
 ### SDXL Turbo #### 
 pipe_turbo = StableDiffusionXLPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
-pipe_turbo.to("cuda")
+#pipe_turbo.to("cuda")
 
 ### SDXL Lightning ### 
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_1step_unet_x0.safetensors" 
 
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16)
 unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
-pipe_lightning = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+pipe_lightning = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")#.to("cuda")
 del unet
 pipe_lightning.scheduler = EulerDiscreteScheduler.from_config(pipe_lightning.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
-pipe_lightning.to("cuda")
+#pipe_lightning.to("cuda")
 
 ### Hyper SDXL ### 
 repo_name = "ByteDance/Hyper-SD"
 ckpt_name = "Hyper-SDXL-1step-Unet.safetensors"
 
-unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name), device="cuda"))
-pipe_hyper = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name))
+pipe_hyper = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16")#.to("cuda")
 pipe_hyper.scheduler = LCMScheduler.from_config(pipe_hyper.scheduler.config)
-pipe_hyper.to("cuda")
+#pipe_hyper.to("cuda")
 del unet
 
 @spaces.GPU
 def run_comparison(prompt):
+    image_turbo.to("cuda")
     image_turbo=pipe_turbo(prompt=prompt, num_inference_steps=1, guidance_scale=0).images[0]
+    image_turbo.to("cpu")
+    image_lightning.to("cuda")
     image_lightning=pipe_lightning(prompt=prompt, num_inference_steps=1, guidance_scale=0).images[0]
+    image_lightning.to("cpu")
+    image_hyper.to("cuda")
     image_hyper=pipe_hyper(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[800]).images[0]
+    image_turbo.to("cpu")
     return image_turbo, image_lightning, image_hyper