Update app.py

app.py

@@ -5,7 +5,6 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
-    TextStreamer,
     set_seed
 )
 from datasets import load_dataset
@@ -44,16 +43,22 @@ doge_generation_config = GenerationConfig(
     repetition_penalty=1.0
 )
 
-# Load speaker embedding with exactly 600 values
+# Load speaker embedding with fallback
 speaker_embedding = None
 embedding_data = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+
 for entry in embedding_data:
     vec = entry["xvector"]
     if len(vec) >= 600:
-        speaker_embedding = torch.tensor(vec[:600], dtype=torch.float32).unsqueeze(0)  # Shape: [1, 600]
+        speaker_embedding = torch.tensor(vec[:600], dtype=torch.float32).unsqueeze(0)
         break
+
+# Fallback: use a zero vector if none found
 if speaker_embedding is None:
-    raise ValueError("No suitable speaker embedding of at least 600 dimensions found.")
+    print("⚠️ No suitable speaker embedding found. Using default 600-dim zero vector.")
+    speaker_embedding = torch.zeros(1, 600, dtype=torch.float32)
+
+# Ensure correct shape
 assert speaker_embedding.shape == (1, 600), f"Expected shape (1, 600), got {speaker_embedding.shape}"
 
 
@@ -75,7 +80,7 @@ def process_image(image):
 
         # 3. Prompt Doge model for context generation
         prompt = f"Determine the context of this image based on the caption and extracted text.\nCaption: {caption}\nExtracted text: {extracted_text}\nContext:"
-        prompt = prompt[:600]  # Ensure prompt isn't too long
+        prompt = prompt[:600]  # Prevent tensor mismatch error
         conversation = [{"role": "user", "content": prompt}]
         doge_inputs = doge_tokenizer.apply_chat_template(
             conversation=conversation,