Update app.py

app.py

@@ -1,34 +1,46 @@
 import gradio as gr
 from transformers import pipeline
 import numpy as np
+from generator import load_csm_1b
+import torchaudio
 
-# Captioning model
+# Load CSM model
+generator = load_csm_1b(device="cpu") 
+
+# Load image-to-text model
 captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Bark TTS model
-synthesiser = pipeline("text-to-speech", model="suno/bark")
+def process_image(input_image):
+    try:
+        # Generate caption
+        caption = captioner(input_image)[0]['generated_text']
+
+        # Generate speech using CSM
+        audio = generator.generate(
+            text=caption,
+            speaker=0,
+            context=[],
+            max_audio_length_ms=10_000,
+        )
 
-def launch(input_image):
-    # Step 1: Generate caption
-    caption = captioner(input_image)[0]['generated_text']
+        # Convert the audio tensor to NumPy for Gradio
+        audio_np = audio.unsqueeze(0).cpu().numpy()
 
-    # Step 2: Synthesize speech from caption
-    speech = synthesiser(caption, forward_params={"do_sample": True})
-    audio = np.array(speech["audio"])
-    rate = speech["sampling_rate"]
+        return (audio_np, generator.sample_rate), caption
 
-    return (audio, rate), caption
+    except Exception as e:
+        return str(e), "Error generating caption or speech."
 
-# Gradio UI
+# Set up Gradio UI
 iface = gr.Interface(
-    fn=launch,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    fn=process_image,
+    inputs=gr.Image(type='pil', label="Upload Image"),
     outputs=[
-        gr.Audio(type="numpy", label="Narrated Audio"),
+        gr.Audio(type="numpy", label="Generated Speech"),
         gr.Textbox(label="Generated Caption")
     ],
-    title="🎙️ SeeSay",
-    description="Upload an image to hear it described with expressive speech."
+    title="🎙️ SeeSay with CSM",
+    description="Upload an image to generate a caption and hear it narrated using CSM."
 )
 
-iface.launch(share = True)
+iface.launch(share=True)