move everything to same device

app.py

@@ -314,10 +314,11 @@ class customUnClipPipeline(UnCLIPImageVariationPipeline):
 
 
 ### ADDITIONAL PIPELINE CODE FOR KARLO
+torch_device = 'cuda'
 pipe = customUnClipPipeline.from_pretrained("kakaobrain/karlo-v1-alpha-image-variations", torch_dtype=torch.float32, trust_remote_code=True, 
-                                            # accelerator='ort', device='cpu'
+                                            # accelerator='ort', device=torch_device
                                            )
-pipe.to('cuda')
+pipe.to(torch_device)
 # pipe.enable_model_cpu_offload()
 
 
@@ -334,11 +335,11 @@ def load_img_from_URL(URL):
 
 def embed_img(input_image):
     tokens = pipe.feature_extractor(input_image)
-    img_model = pipe.image_encoder.to('cpu')
+    img_model = pipe.image_encoder.to(torch_device)
     with torch.no_grad():
         embeds = img_model(torch.tensor(tokens.pixel_values[0]).unsqueeze(0))
         
-    return embeds.image_embeds.to('cpu')
+    return embeds.image_embeds.to(torch_device)
 
 
 def localimg_2_embed(image_dir):
@@ -389,16 +390,16 @@ def image_grid(imgs, rows, cols):
     return grid
 
 
-chaosclicker_willtensor = localimg_2_embed('willpaint-imgs/chaosclicker-willpaint.png').to('cpu')
-contentcnsr_willtensor = localimg_2_embed('willpaint-imgs/contentconnoisseur-willpaint.png').to('cpu')
-digdaydrmr_willtensor = localimg_2_embed('willpaint-imgs/digitaldaydreamer-willpaint.png').to('cpu')
-ecoexplr_willtensor = localimg_2_embed('willpaint-imgs/ecoexplorer-willpaint.png').to('cpu')
-fandomfox_willtensor = localimg_2_embed('willpaint-imgs/fandomfox-willpaint.png').to('cpu')
-mememaven_willtensor = localimg_2_embed('willpaint-imgs/mememaven-willpaint.png').to('cpu')
-newsnerd_willtensor = localimg_2_embed('willpaint-imgs/newnerd-willpaint.png').to('cpu')
-nostalgicnvgtr_willtensor = localimg_2_embed('willpaint-imgs/nostalgicnavigator-willpaint.png').to('cpu')
-scrollseeker_willtensor = localimg_2_embed('willpaint-imgs/scrollseeker-willpaint.png').to('cpu')
-trendtracker_willtensor = localimg_2_embed('willpaint-imgs/trendtracker-willpaint.png').to('cpu')
+chaosclicker_willtensor = localimg_2_embed('willpaint-imgs/chaosclicker-willpaint.png').to(torch_device)
+contentcnsr_willtensor = localimg_2_embed('willpaint-imgs/contentconnoisseur-willpaint.png').to(torch_device)
+digdaydrmr_willtensor = localimg_2_embed('willpaint-imgs/digitaldaydreamer-willpaint.png').to(torch_device)
+ecoexplr_willtensor = localimg_2_embed('willpaint-imgs/ecoexplorer-willpaint.png').to(torch_device)
+fandomfox_willtensor = localimg_2_embed('willpaint-imgs/fandomfox-willpaint.png').to(torch_device)
+mememaven_willtensor = localimg_2_embed('willpaint-imgs/mememaven-willpaint.png').to(torch_device)
+newsnerd_willtensor = localimg_2_embed('willpaint-imgs/newnerd-willpaint.png').to(torch_device)
+nostalgicnvgtr_willtensor = localimg_2_embed('willpaint-imgs/nostalgicnavigator-willpaint.png').to(torch_device)
+scrollseeker_willtensor = localimg_2_embed('willpaint-imgs/scrollseeker-willpaint.png').to(torch_device)
+trendtracker_willtensor = localimg_2_embed('willpaint-imgs/trendtracker-willpaint.png').to(torch_device)
 
 
 will_cand_tensors = torch.cat([chaosclicker_willtensor,