fix image generation

app.py

@@ -5,7 +5,7 @@ import os
 import spaces
 import torch
 from diffusers import DiffusionPipeline
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 from huggingface_hub import login
 
 hf_token = os.getenv("hf_token")
@@ -37,17 +37,22 @@ def get_image_gen_pipeline():
 def get_text_gen_pipeline():
     global _text_gen_pipeline
     if _text_gen_pipeline is None:
-         try:
-              device = "cuda" if torch.cuda.is_available() else "cpu"
-              _text_gen_pipeline = pipeline(
-                   "text-generation",
-                   model="mistralai/Mistral-7B-Instruct-v0.3",
-                   max_new_tokens=2048,
-                   device=device,
-              )
-         except Exception as e:
-              print(f"Error loading text generation model: {e}")
-              return None
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            tokenizer = AutoTokenizer.from_pretrained(
+                "mistralai/Mistral-7B-Instruct-v0.3",
+                use_fast=True  # Ensure fast tokenizer
+            )
+            _text_gen_pipeline = pipeline(
+                "text-generation",
+                model="mistralai/Mistral-7B-Instruct-v0.3",
+                tokenizer=tokenizer,
+                max_new_tokens=2048,
+                device=device,
+            )
+        except Exception as e:
+            print(f"Error loading text generation model: {e}")
+            return None
     return _text_gen_pipeline
 
 @spaces.GPU()