Update app.py

app.py

@@ -63,7 +63,6 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
         elif isinstance(caption, str):
             captions.append(caption)
         elif isinstance(caption, (list, np.ndarray)):
-            # take a random caption if there are multiple
             captions.append(random.choice(caption) if is_train else caption[0])
 
     with torch.no_grad():
@@ -76,12 +75,14 @@ def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompt
             return_tensors="pt",
         )
 
-        text_input_ids = text_inputs.input_ids
-        prompt_masks = text_inputs.attention_mask
+    
+        device = text_encoder.device
+        text_input_ids = text_inputs.input_ids.to(device)
+        prompt_masks = text_inputs.attention_mask.to(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]