[AuraFlow] fix long prompt handling (#8937)

sayakpaul · sayakpaul · commit 48ad9a9e9302 · 2024-12-23T13:02:14.000+05:30
fix
diff --git a/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py b/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
@@ -260,7 +260,6 @@ def encode_prompt(
                 padding="max_length",
                 return_tensors="pt",
             )
-            text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
             text_input_ids = text_inputs["input_ids"]
             untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
 
@@ -273,6 +272,7 @@ def encode_prompt(
                     f" {max_length} tokens: {removed_text}"
                 )
 
+            text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
             prompt_embeds = self.text_encoder(**text_inputs)[0]
             prompt_attention_mask = text_inputs["attention_mask"].unsqueeze(-1).expand(prompt_embeds.shape)
             prompt_embeds = prompt_embeds * prompt_attention_mask

Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,6 @@ def encode_prompt(`
`260`	`260`	`padding="max_length",`
`261`	`261`	`return_tensors="pt",`
`262`	`262`	`)`
`263`		`- text_inputs = {k: v.to(device) for k, v in text_inputs.items()}`
`264`	`263`	`text_input_ids = text_inputs["input_ids"]`
`265`	`264`	`untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids`
`266`	`265`
`@@ -273,6 +272,7 @@ def encode_prompt(`
`273`	`272`	`f" {max_length} tokens: {removed_text}"`
`274`	`273`	`)`
`275`	`274`
	`275`	`+ text_inputs = {k: v.to(device) for k, v in text_inputs.items()}`
`276`	`276`	`prompt_embeds = self.text_encoder(**text_inputs)[0]`
`277`	`277`	`prompt_attention_mask = text_inputs["attention_mask"].unsqueeze(-1).expand(prompt_embeds.shape)`
`278`	`278`	`prompt_embeds = prompt_embeds * prompt_attention_mask`