device changes

app.py

@@ -34,18 +34,18 @@ 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").to("cuda", torch.float16)
-unet.load_state_dict(load_file(hf_hub_download(sdxl_lightening, ckpt), device="cuda"))
+unet = UNet2DConditionModel.from_config(model_id, subfolder="unet").to(DEVICE, torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(sdxl_lightening, ckpt), device=DEVICE))
 
-image_encoder = CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter",  subfolder="models/image_encoder", torch_dtype=torch.float16,).to("cuda")
-pipe = AutoPipelineForText2Image.from_pretrained(model_id, unet=unet, torch_dtype=torch.float16, variant="fp16", image_encoder=image_encoder).to("cuda")
+image_encoder = CLIPVisionModelWithProjection.from_pretrained("h94/IP-Adapter",  subfolder="models/image_encoder", torch_dtype=torch.float16,).to(DEVICE)
+pipe = AutoPipelineForText2Image.from_pretrained(model_id, unet=unet, torch_dtype=torch.float16, variant="fp16", image_encoder=image_encoder).to(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.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-pipe.to(device='cuda')
+pipe.to(device=DEVICE)
 
 
 output_hidden_state = False
@@ -60,13 +60,13 @@ def predict(
     """Run a single prediction on the model"""
     with torch.no_grad():
         if im_emb == None:
-            im_emb = torch.zeros(1, 1024, dtype=torch.float16, device='cuda')
+            im_emb = torch.zeros(1, 1024, dtype=torch.float16, device=DEVICE)
             
-        im_emb = [im_emb.to('cuda').unsqueeze(0)]
+        im_emb = [im_emb.to(DEVICE).unsqueeze(0)]
         if prompt == '':
             image = pipe(
-                prompt_embeds=torch.zeros(1, 1, 2048, dtype=torch.float16, device='cuda'),
-                pooled_prompt_embeds=torch.zeros(1, 1280, dtype=torch.float16, device='cuda'),
+                prompt_embeds=torch.zeros(1, 1, 2048, dtype=torch.float16, device=DEVICE),
+                pooled_prompt_embeds=torch.zeros(1, 1280, dtype=torch.float16, device=DEVICE),
                 ip_adapter_image_embeds=im_emb,
                 height=1024,
                 width=1024,                                                                                                             
@@ -83,9 +83,9 @@ def predict(
                 guidance_scale=0,
             ).images[0]
         im_emb, _ = pipe.encode_image(
-                image, 'cuda', 1, output_hidden_state
+                image, DEVICE, 1, output_hidden_state
             )
-        return image, im_emb.to(DEVICE)
+        return image, im_emb.to('cpu')
 
 # TODO add to state instead of shared across all
 glob_idx = 0
@@ -128,7 +128,7 @@ def next_image(embs, ys, calibrate_prompts):
                 if has_0 and has_1:
                     break
             
-            feature_embs = np.array(torch.cat([embs[i] for i in indices]).to('cpu'))
+            feature_embs = np.array(torch.cat([embs[i].to('cpu') for i in indices]).to('cpu'))
             scaler = preprocessing.StandardScaler().fit(feature_embs)
             feature_embs = scaler.transform(feature_embs)
 
@@ -138,7 +138,7 @@ def next_image(embs, ys, calibrate_prompts):
 
             rng_prompt = random.choice(prompt_list)
             w = 1# if len(embs) % 2 == 0 else 0
-            im_emb = w * lin_class.coef_.to(device=DEVICE, dtype=torch.float16)
+            im_emb = w * lin_class.coef_.to(dtype=torch.float16)
             prompt= 'an image' if glob_idx % 2 == 0 else rng_prompt
             print(prompt, len(ys))
             image, im_emb = predict(prompt, im_emb)