Update app.py

Add lazy loading for models to prevent overuse of resources.

app.py

@@ -51,12 +51,15 @@ 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")
+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)
 
 # Load speaker embeddings for text-to-audio
 def load_speaker_embeddings(model_name):
@@ -64,26 +67,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).to()  # 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")
+    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:
         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()
+        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]
@@ -92,7 +95,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()
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(device)
     
     # Generate a response
     with torch.no_grad():
@@ -112,8 +115,8 @@ 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, _map="auto"
-        )
+            text_to_image_models[model_name], use_auth_token=read_token, torch_dtype=torch.float16, device_map="auto"
+        ).to(device)
     pipeline = text_to_image_pipelines[model_name]
     image = pipeline(prompt).images[0]
     return image
@@ -122,7 +125,7 @@ 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, speaker_embeddings=speaker_embeddings)