Update app.py

app.py

@@ -76,12 +76,14 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
         )
 
         print(f"Text Encoder Device: {text_encoder.device}")
-        text_input_ids = text_inputs.input_ids
-        prompt_masks = text_inputs.attention_mask
+        text_input_ids = text_inputs.input_ids.cuda()
+        prompt_masks = text_inputs.attention_mask.cuda()
+        print(f"Text Input Ids Device: {text_input_ids.device}")
+        print(f"Prompt Masks Device: {prompt_masks.device}")
 
         prompt_embeds = text_encoder(
-            input_ids=text_input_ids.cuda(),
-            attention_mask=prompt_masks.cuda(),
+            input_ids=text_input_ids,
+            attention_mask=prompt_masks,
             output_hidden_states=True,
         ).hidden_states[-2]