Update app.py

app.py

@@ -16,6 +16,10 @@ if not read_token:
 from huggingface_hub import login
 login(read_token)
 
+# Set device to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logger.info(f"Device set to use {device}")
+
 # Define a dictionary of conversational models
 conversational_models = {
     "Qwen": "Qwen/QwQ-32B",
@@ -49,12 +53,12 @@ text_to_image_pipelines = {}
 text_to_speech_pipelines = {}
 
 # Initialize pipelines for other tasks
-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):
@@ -62,26 +66,26 @@ 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)  # Example speaker
+        speaker_embeddings = torch.tensor(dataset[7306]["xvector"]).unsqueeze(0).to(device)  # 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", use_auth_token=read_token)
+    text_to_audio_pipeline = pipeline("text-to-audio", model="stabilityai/stable-audio-open-1.0", use_auth_token=read_token, 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", use_auth_token=read_token, 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:
         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)
+        model = AutoModelForCausalLM.from_pretrained(conversational_models[model_name], use_auth_token=read_token).to(device)
         conversational_tokenizers[model_name] = tokenizer
         conversational_models_loaded[model_name] = model
     return conversational_tokenizers[model_name], conversational_models_loaded[model_name]
@@ -90,7 +94,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():
@@ -110,7 +114,7 @@ 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
+            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]
@@ -120,12 +124,14 @@ 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
+            "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)
+    audio = pipeline(text, speaker_embeddings=speaker_embeddings)
     return audio["audio"]
 
+
+
 def visual_qa(image, question):
     result = visual_qa_pipeline(image, question)
     return result["answer"]