Update app.py

app.py

@@ -32,22 +32,12 @@ text_to_speech_models = {
 conversational_tokenizers = {}
 conversational_models_loaded = {}
 
-for model_name, model_id in conversational_models.items():
-    conversational_tokenizers[model_name] = AutoTokenizer.from_pretrained(model_id)
-    conversational_models_loaded[model_name] = AutoModelForCausalLM.from_pretrained(model_id)
-
 # Initialize pipelines for Text-to-Image
 text_to_image_pipelines = {}
 
-for model_name, model_id in text_to_image_models.items():
-    text_to_image_pipelines[model_name] = StableDiffusionPipeline.from_pretrained(model_id)
-
 # Initialize pipelines for Text-to-Speech
 text_to_speech_pipelines = {}
 
-for model_name, model_id in text_to_speech_models.items():
-    text_to_speech_pipelines[model_name] = pipeline("text-to-speech", model=model_id)
-
 # 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")
@@ -61,9 +51,16 @@ summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cn
 text_to_audio_pipeline = pipeline("text-to-speech", model="julien-c/ljspeech_tts_train_tacotron2_raw_phn_tacotron_g2p_en_no_space")
 audio_classification_pipeline = pipeline("audio-classification", model="facebook/wav2vec2-base")
 
+def load_conversational_model(model_name):
+    if model_name not in conversational_models_loaded:
+        tokenizer = AutoTokenizer.from_pretrained(conversational_models[model_name])
+        model = AutoModelForCausalLM.from_pretrained(conversational_models[model_name])
+        conversational_tokenizers[model_name] = tokenizer
+        conversational_models_loaded[model_name] = model
+    return conversational_tokenizers[model_name], conversational_models_loaded[model_name]
+
 def chat(model_name, user_input, history=[]):
-    tokenizer = conversational_tokenizers[model_name]
-    model = conversational_models_loaded[model_name]
+    tokenizer, model = load_conversational_model(model_name)
     
     # Encode the input
     input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
@@ -83,11 +80,15 @@ def chat(model_name, user_input, history=[]):
     return history, history
 
 def generate_image(model_name, prompt):
+    if model_name not in text_to_image_pipelines:
+        text_to_image_pipelines[model_name] = StableDiffusionPipeline.from_pretrained(text_to_image_models[model_name])
     pipeline = text_to_image_pipelines[model_name]
     image = pipeline(prompt).images[0]
     return image
 
 def generate_speech(model_name, text):
+    if model_name not in text_to_speech_pipelines:
+        text_to_speech_pipelines[model_name] = pipeline("text-to-speech", model=text_to_speech_models[model_name])
     pipeline = text_to_speech_pipelines[model_name]
     audio = pipeline(text)
     return audio["audio"]
@@ -235,14 +236,4 @@ with gr.Blocks() as demo:
         text_to_audio_generate = gr.Button("Generate Audio")
         text_to_audio_output = gr.Audio(label="Generated Audio")
         
-        text_to_audio_generate.click(text_to_audio, inputs=text_to_audio_text, outputs=text_to_audio_output)
-    
-    with gr.Tab("Audio Classification"):
-        audio_classification_audio = gr.Audio(label="Upload Audio")
-        audio_classification_generate = gr.Button("Classify")
-        audio_classification_output = gr.Textbox(label="Classification Result")
-        
-        audio_classification_generate.click(audio_classification, inputs=audio_classification_audio, outputs=audio_classification_output)
-
-# Launch the demo
-demo.launch()
+        text_to_audio_generate.click(text_to_audio, inputs=text_to_audio_text, outputs=text_to_audio_output)