Skip to content

Commit b351be2

Browse files
authored
LongCat Image pipeline: Allow offloading/quantization of text_encoder component (#12963)
* Don't attempt to move the text_encoder. Just move the generated_ids. * The inputs to the text_encoder should be on its device
1 parent d8f4dd2 commit b351be2

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,10 +260,10 @@ def rewire_prompt(self, prompt, device):
260260
text = self.text_processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
261261
all_text.append(text)
262262

263-
inputs = self.text_processor(text=all_text, padding=True, return_tensors="pt").to(device)
263+
inputs = self.text_processor(text=all_text, padding=True, return_tensors="pt").to(self.text_encoder.device)
264264

265-
self.text_encoder.to(device)
266265
generated_ids = self.text_encoder.generate(**inputs, max_new_tokens=self.tokenizer_max_length)
266+
generated_ids.to(device)
267267
generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
268268
output_text = self.text_processor.batch_decode(
269269
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False

0 commit comments

Comments
 (0)