mpatel57 commited on
Commit
b75436e
1 Parent(s): d340330

device execution changed

Browse files
src/pipelines/pipeline_kandinsky_subject_prior.py CHANGED
@@ -469,7 +469,7 @@ class KandinskyPriorPipeline(DiffusionPipeline):
469
  truncation=True,
470
  return_tensors="pt",
471
  )
472
- txt_items = {k: v.to("cuda") for k, v in txt.items()}
473
  txt_feats = self.text_encoder(**txt_items)
474
  last_hidden_states = txt_feats.last_hidden_state[0].detach().cpu().numpy()
475
  prompt_embeds = txt_feats.text_embeds.detach().cpu()
@@ -478,7 +478,7 @@ class KandinskyPriorPipeline(DiffusionPipeline):
478
  for sub_img, sub_name in zip(raw_data["subject_images"], raw_data["subject_keywords"]):
479
  if isinstance(sub_img, str):
480
  sub_img = Image.open(sub_img)
481
- mask_img = self.image_processor(sub_img, return_tensors="pt").to("cuda")
482
  vision_feats = self.image_encoder(**mask_img).image_embeds
483
  entity_tokens = self.tokenizer(sub_name)["input_ids"][1:-1]
484
 
 
469
  truncation=True,
470
  return_tensors="pt",
471
  )
472
+ txt_items = {k: v.to(self._execution_device) for k, v in txt.items()}
473
  txt_feats = self.text_encoder(**txt_items)
474
  last_hidden_states = txt_feats.last_hidden_state[0].detach().cpu().numpy()
475
  prompt_embeds = txt_feats.text_embeds.detach().cpu()
 
478
  for sub_img, sub_name in zip(raw_data["subject_images"], raw_data["subject_keywords"]):
479
  if isinstance(sub_img, str):
480
  sub_img = Image.open(sub_img)
481
+ mask_img = self.image_processor(sub_img, return_tensors="pt").to(self._execution_device)
482
  vision_feats = self.image_encoder(**mask_img).image_embeds
483
  entity_tokens = self.tokenizer(sub_name)["input_ids"][1:-1]
484