.

app.py

@@ -39,17 +39,17 @@ start_time = time.time()
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 sdxl_lightening = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_2step_unet.safetensors"
-unet = UNet2DConditionModel.from_config(model_id, subfolder="unet", low_cpu_mem_usage=True).to(torch.float16)
+unet = UNet2DConditionModel.from_config(model_id, subfolder="unet", low_cpu_mem_usage=True, device=DEVICE).to(torch.float16)
 unet.load_state_dict(load_file(hf_hub_download(sdxl_lightening, ckpt)))
 
-image_encoder = CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter",  subfolder="models/image_encoder", torch_dtype=torch.float16, low_cpu_mem_usage=True)
-pipe = AutoPipelineForText2Image.from_pretrained(model_id, unet=unet, torch_dtype=torch.float16, variant="fp16", image_encoder=image_encoder, low_cpu_mem_usage=True)
+image_encoder = CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter",  subfolder="models/image_encoder", torch_dtype=torch.float16, low_cpu_mem_usage=True, device=DEVICE)
+pipe = AutoPipelineForText2Image.from_pretrained(model_id, unet=unet, torch_dtype=torch.float16, variant="fp16", image_encoder=image_encoder, low_cpu_mem_usage=True, device=DEVICE)
 pipe.unet._load_ip_adapter_weights(torch.load(hf_hub_download('h94/IP-Adapter', 'sdxl_models/ip-adapter_sdxl_vit-h.bin')))
 pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl_vit-h.bin")
 pipe.register_modules(image_encoder = image_encoder)
 pipe.set_ip_adapter_scale(0.8)
 
-pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16, low_cpu_mem_usage=True, device=DEVICE)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
 pipe.to(device=DEVICE)
@@ -75,7 +75,7 @@ class BottleneckT5Autoencoder:
     def __init__(self, model_path: str, device='cuda'):
         self.device = device
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512, torch_dtype=torch.bfloat16)
-        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(self.device)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device=DEVICE, low_cpu_mem_usage=True)
         self.model.eval()