check te device

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-10-14 09:29:33 -04:00
parent a2c5893f7f
commit c2ab0b11c3
3 changed files with 23 additions and 14 deletions
+11 -2
View File
@@ -18,6 +18,8 @@ cache = {}
def compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor:
if not devices.same_device(self.text_encoder.device, devices.device):
sd_models.move_model(self.text_encoder, devices.device)
needs_hidden_states = self.returned_embeddings_type != 1
text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True)
@@ -298,13 +300,20 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]:
no_mask_provider = EmbeddingsProvider(padding_attention_mask_value=1 if "sote" in pipe.sd_checkpoint_info.name.lower() else 0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(no_mask_provider)
elif getattr(pipe, "tokenizer", None) is not None and getattr(pipe, "text_encoder", None) is not None:
sd_models.move_model(pipe.text_encoder, device)
if not devices.same_device(pipe.text_encoder.device, devices.device):
sd_models.move_model(pipe.text_encoder, devices.device)
provider = EmbeddingsProvider(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(provider)
if getattr(pipe, "tokenizer_2", None) is not None and getattr(pipe, "text_encoder_2", None) is not None:
sd_models.move_model(pipe.text_encoder, device)
if not devices.same_device(pipe.text_encoder_2.device, devices.device):
sd_models.move_model(pipe.text_encoder_2, devices.device)
provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_2, text_encoder=pipe.text_encoder_2, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(provider)
if getattr(pipe, "tokenizer_3", None) is not None and getattr(pipe, "text_encoder_3", None) is not None:
if not devices.same_device(pipe.text_encoder_3.device, devices.device):
sd_models.move_model(pipe.text_encoder_3, devices.device)
provider = EmbeddingsProvider(tokenizer=pipe.tokenizer_3, text_encoder=pipe.text_encoder_3, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(provider)
return embeddings_providers