Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusion3Pipeline, StableDiffusionPipeline, DiffusionPipeline, StableDiffusionXLPipeline, DPMSolverSinglestepScheduler
+from diffusers import StableDiffusion3Pipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverSinglestepScheduler
 import gradio as gr
 import os
 import random
@@ -22,22 +22,22 @@ MAX_SEED = np.iinfo(np.int32).max
 sd3_medium_pipe = StableDiffusion3Pipeline.from_pretrained(
     "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
 )
-sd3_medium_pipe.to(device)
+sd3_medium_pipe.enable_model_cpu_offload()
 
 sd2_1_pipe = StableDiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16
 )
-sd2_1_pipe.to(device)
+sd2_1_pipe.enable_model_cpu_offload()
 
-sdxl_pipe = DiffusionPipeline.from_pretrained(
+sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
 )
-sdxl_pipe.to(device)
+sdxl_pipe.enable_model_cpu_offload()
 
 sdxl_flash_pipe = StableDiffusionXLPipeline.from_pretrained(
     "sd-community/sdxl-flash", torch_dtype=torch.float16
 )
-sdxl_flash_pipe.to(device)
+sdxl_flash_pipe.enable_model_cpu_offload()
 # Ensure sampler uses "trailing" timesteps for sdxl flash.
 sdxl_flash_pipe.scheduler = DPMSolverSinglestepScheduler.from_config(sdxl_flash_pipe.scheduler.config, timestep_spacing="trailing")