Update app.py

app.py

@@ -54,12 +54,12 @@ text_to_speech_pipelines = {}
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Device set to use {device}")
 
-visual_qa_pipeline = pipeline("visual-question-answering", model="dandelin/vilt-b32-finetuned-vqa", device=device)
-document_qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2", device=device)
-image_classification_pipeline = pipeline("image-classification", model="facebook/deit-base-distilled-patch16-224", device=device)
-object_detection_pipeline = pipeline("object-detection", model="facebook/detr-resnet-50", device=device)
-video_classification_pipeline = pipeline("video-classification", model="facebook/timesformer-base-finetuned-k400", device=device)
-summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cnn", device=device)
+visual_qa_pipeline = pipeline("visual-question-answering", model="dandelin/vilt-b32-finetuned-vqa")
+document_qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
+image_classification_pipeline = pipeline("image-classification", model="facebook/deit-base-distilled-patch16-224")
+object_detection_pipeline = pipeline("object-detection", model="facebook/detr-resnet-50")
+video_classification_pipeline = pipeline("video-classification", model="facebook/timesformer-base-finetuned-k400")
+summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
 
 # Load speaker embeddings for text-to-audio
 def load_speaker_embeddings(model_name):
@@ -67,26 +67,34 @@ def load_speaker_embeddings(model_name):
         logger.info("Loading speaker embeddings for SpeechT5")
         from datasets import load_dataset
         dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-        speaker_embeddings = torch.tensor(dataset[7306]["xvector"]).unsqueeze(0).to(device)  # Example speaker
+        speaker_embeddings = torch.tensor(dataset[7306]["xvector"]).unsqueeze(0)  # Example speaker
         return speaker_embeddings
     return None
 
 # Use a different model for text-to-audio if stabilityai/stable-audio-open-1.0 is not supported
 try:
-    text_to_audio_pipeline = pipeline("text-to-audio", model="stabilityai/stable-audio-open-1.0", device=device)
+    text_to_audio_pipeline = pipeline("text-to-audio", model="stabilityai/stable-audio-open-1.0")
 except ValueError as e:
     logger.error(f"Error loading stabilityai/stable-audio-open-1.0: {e}")
     logger.info("Falling back to a different text-to-audio model.")
-    text_to_audio_pipeline = pipeline("text-to-audio", model="microsoft/speecht5_tts", device=device)
+    text_to_audio_pipeline = pipeline("text-to-audio", model="microsoft/speecht5_tts")
     speaker_embeddings = load_speaker_embeddings("microsoft/speecht5_tts")
 
-audio_classification_pipeline = pipeline("audio-classification", model="facebook/wav2vec2-base", device=device)
+audio_classification_pipeline = pipeline("audio-classification", model="facebook/wav2vec2-base")
 
 def load_conversational_model(model_name):
     if model_name not in conversational_models_loaded:
         logger.info(f"Loading conversational model: {model_name}")
-        tokenizer = AutoTokenizer.from_pretrained(conversational_models[model_name], use_auth_token=read_token)
-        model = AutoModelForCausalLM.from_pretrained(conversational_models[model_name], use_auth_token=read_token).to(device)
+        tokenizer = AutoTokenizer.from_pretrained(
+            conversational_models[model_name], 
+            use_auth_token=read_token, 
+            trust_remote_code=True
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            conversational_models[model_name], 
+            use_auth_token=read_token, 
+            trust_remote_code=True
+        )
         conversational_tokenizers[model_name] = tokenizer
         conversational_models_loaded[model_name] = model
     return conversational_tokenizers[model_name], conversational_models_loaded[model_name]
@@ -95,7 +103,7 @@ def chat(model_name, user_input, history=[]):
     tokenizer, model = load_conversational_model(model_name)
     
     # Encode the input
-    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(device)
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
     
     # Generate a response
     with torch.no_grad():
@@ -115,8 +123,11 @@ def generate_image(model_name, prompt):
     if model_name not in text_to_image_pipelines:
         logger.info(f"Loading text-to-image model: {model_name}")
         text_to_image_pipelines[model_name] = StableDiffusionPipeline.from_pretrained(
-            text_to_image_models[model_name], use_auth_token=read_token, torch_dtype=torch.float16, device_map="auto"
-        ).to(device)
+            text_to_image_models[model_name], 
+            use_auth_token=read_token, 
+            torch_dtype=torch.float16, 
+            device_map="auto"
+        )
     pipeline = text_to_image_pipelines[model_name]
     image = pipeline(prompt).images[0]
     return image
@@ -125,7 +136,10 @@ def generate_speech(model_name, text):
     if model_name not in text_to_speech_pipelines:
         logger.info(f"Loading text-to-speech model: {model_name}")
         text_to_speech_pipelines[model_name] = pipeline(
-            "text-to-speech", model=text_to_speech_models[model_name], use_auth_token=read_token, device=device
+            "text-to-speech", 
+            model=text_to_speech_models[model_name], 
+            use_auth_token=read_token, 
+            device=device
         )
     pipeline = text_to_speech_pipelines[model_name]
     audio = pipeline(text, speaker_embeddings=speaker_embeddings)