Update app.py

app.py

@@ -4,8 +4,7 @@ from diffusers import StableDiffusionPipeline
 import torch
 import os
 import logging
-from huggingface_hub import login, hf_hub_download
-import requests
+from huggingface_hub import login
 import accelerate
 
 # Set up logging
@@ -56,12 +55,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")
-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")
+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)
 
 # Load speaker embeddings for text-to-audio
 def load_speaker_embeddings(model_name):
@@ -75,14 +74,14 @@ def load_speaker_embeddings(model_name):
 
 # 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")
+    text_to_audio_pipeline = pipeline("text-to-audio", model="stabilityai/stable-audio-open-1.0", device=device)
 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")
+    text_to_audio_pipeline = pipeline("text-to-audio", model="microsoft/speecht5_tts", device=device)
     speaker_embeddings = load_speaker_embeddings("microsoft/speecht5_tts")
 
-audio_classification_pipeline = pipeline("audio-classification", model="facebook/wav2vec2-base")
+audio_classification_pipeline = pipeline("audio-classification", model="facebook/wav2vec2-base", device=device)
 
 def load_conversational_model(model_name):
     if model_name not in conversational_models_loaded:
@@ -92,11 +91,22 @@ def load_conversational_model(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
-        )
+        try:
+            model = AutoModelForCausalLM.from_pretrained(
+                conversational_models[model_name], 
+                use_auth_token=read_token, 
+                trust_remote_code=True,
+                device_map="auto" if torch.cuda.is_available() else "cpu"
+            )
+        except RuntimeError as e:
+            logger.error(f"RuntimeError: {e}")
+            logger.info("Falling back to CPU for the model.")
+            model = AutoModelForCausalLM.from_pretrained(
+                conversational_models[model_name], 
+                use_auth_token=read_token, 
+                trust_remote_code=True,
+                device_map="cpu"
+            )
         conversational_tokenizers[model_name] = tokenizer
         conversational_models_loaded[model_name] = model
     return conversational_tokenizers[model_name], conversational_models_loaded[model_name]
@@ -105,7 +115,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")
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(device)
     
     # Generate a response
     with torch.no_grad():
@@ -127,8 +137,8 @@ def generate_image(model_name, prompt):
         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"
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, 
+            device_map="auto" if torch.cuda.is_available() else "cpu"
         )
     pipeline = text_to_image_pipelines[model_name]
     image = pipeline(prompt).images[0]