fix tensor mismatch

app.py

@@ -4,6 +4,7 @@ import torch
 from diffusers import DiffusionPipeline
 from transformers import T5EncoderModel
 import tempfile
+import os
 
 # Global variable to store the text pipeline
 text_pipe = None
@@ -15,7 +16,6 @@ def load_model():
         print("Loading T5 text encoder...")
         
         # Get token from environment
-        import os
         token = os.getenv("HF_TOKEN")
         
         text_encoder = T5EncoderModel.from_pretrained(
@@ -24,13 +24,14 @@ def load_model():
             load_in_8bit=True,
             variant="8bit",
             device_map="auto",
-            token=token  # Add this line
+            token=token
         )
         text_pipe = DiffusionPipeline.from_pretrained(
             "DeepFloyd/IF-I-L-v1.0",
             text_encoder=text_encoder,
             unet=None,
-            token=token  # Add this line
+            token=token,
+            device_map="auto"  # Add this
         )
         print("Model loaded successfully!")
     return text_pipe
@@ -48,6 +49,12 @@ def generate_embeddings(prompts_text):
         # Load model if not already loaded
         pipe = load_model()
         
+        # Move pipeline to CUDA if available
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+            if hasattr(pipe, 'text_encoder') and pipe.text_encoder is not None:
+                pipe.text_encoder = pipe.text_encoder.to(device)
+        
         # Parse prompts (one per line)
         prompts = [p.strip() for p in prompts_text.strip().split('\n') if p.strip()]
         
@@ -68,8 +75,11 @@ def generate_embeddings(prompts_text):
         # Extract positive prompt embeddings
         prompt_embeds, negative_prompt_embeds = zip(*prompt_embeds_list)
         
+        # Move embeddings to CPU before saving
+        prompt_embeds_cpu = [emb.cpu() if isinstance(emb, torch.Tensor) else emb for emb in prompt_embeds]
+        
         # Create dictionary
-        prompt_embeds_dict = dict(zip(prompts, prompt_embeds))
+        prompt_embeds_dict = dict(zip(prompts, prompt_embeds_cpu))
         
         # Save to temporary file
         temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.pth')
@@ -83,7 +93,9 @@ def generate_embeddings(prompts_text):
         return temp_file.name, status_msg
         
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        import traceback
+        error_details = traceback.format_exc()
+        return None, f"❌ Error: {str(e)}\n\nDetails:\n{error_details}"
 
 # Create Gradio interface
 with gr.Blocks(title="T5 Text Encoder - Embeddings Generator") as demo: