Update app.py

app.py

@@ -1,182 +1,137 @@
-import gradio as gr
 import torch
-import spaces
-from transformers import AutoModel, AutoProcessor
-
-from gender_classification import gender_classification
-from emotion_classification import emotion_classification
-from dog_breed import dog_breed_classification
-from deepfake_quality import deepfake_classification
-from gym_workout_classification import workout_classification
-from augmented_waste_classifier import waste_classification
-from age_classification import age_classification
-from mnist_digits import classify_digit
-from fashion_mnist_cloth import fashion_mnist_classification
-from indian_western_food_classify import food_classification
-from bird_species import bird_classification
-from alphabet_sign_language_detection import sign_language_classification
-from rice_leaf_disease import classify_leaf_disease
-from traffic_density import traffic_density_classification
-from clip_art import clipart_classification
-from multisource_121 import multisource_classification
-from painting_126 import painting_classification
-from sketch_126 import sketch_classification  # New import
-
-# Main classification function for multi-model classification.
-def classify(image, model_name):
-    if model_name == "gender":
-        return gender_classification(image)
-    elif model_name == "emotion":
-        return emotion_classification(image)
-    elif model_name == "dog breed":
-        return dog_breed_classification(image)
-    elif model_name == "deepfake":
-        return deepfake_classification(image)
-    elif model_name == "gym workout":
-        return workout_classification(image)
-    elif model_name == "waste":
-        return waste_classification(image)
-    elif model_name == "age":
-        return age_classification(image)
-    elif model_name == "mnist":
-        return classify_digit(image)
-    elif model_name == "fashion_mnist":
-        return fashion_mnist_classification(image)
-    elif model_name == "food":
-        return food_classification(image)
-    elif model_name == "bird":
-        return bird_classification(image)
-    elif model_name == "leaf disease":
-        return classify_leaf_disease(image)
-    elif model_name == "sign language":
-        return sign_language_classification(image)
-    elif model_name == "traffic density":
-        return traffic_density_classification(image)
-    elif model_name == "clip art":
-        return clipart_classification(image)
-    elif model_name == "multisource":
-        return multisource_classification(image)
-    elif model_name == "painting":
-        return painting_classification(image)
-    elif model_name == "sketch":  # New option
-        return sketch_classification(image)
-    else:
-        return {"Error": "No model selected"}
-
-# Function to update the selected model and button styles.
-def select_model(model_name):
-    model_variants = {
-        "gender": "secondary", "emotion": "secondary", "dog breed": "secondary", "deepfake": "secondary",
-        "gym workout": "secondary", "waste": "secondary", "age": "secondary", "mnist": "secondary",
-        "fashion_mnist": "secondary", "food": "secondary", "bird": "secondary", "leaf disease": "secondary",
-        "sign language": "secondary", "traffic density": "secondary", "clip art": "secondary",
-        "multisource": "secondary", "painting": "secondary", "sketch": "secondary"  # New model variant
-    }
-    model_variants[model_name] = "primary"
-    return (model_name, *(gr.update(variant=model_variants[key]) for key in model_variants))
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+from snac import SNAC
 
-# Zero-Shot Classification Setup (SigLIP models)
-sg1_ckpt = "google/siglip-so400m-patch14-384"
-siglip1_model = AutoModel.from_pretrained(sg1_ckpt, device_map="cpu").eval()
-siglip1_processor = AutoProcessor.from_pretrained(sg1_ckpt)
+def redistribute_codes(row):
+    """
+    Convert a sequence of token codes into an audio waveform using SNAC.
+    The code assumes each 7 tokens represent one group of instructions.
+    """
+    row_length = row.size(0)
+    new_length = (row_length // 7) * 7
+    trimmed_row = row[:new_length]
+    code_list = [t - 128266 for t in trimmed_row]
+    
+    layer_1, layer_2, layer_3 = [], [], []
+    
+    for i in range((len(code_list) + 1) // 7):
+        layer_1.append(code_list[7 * i][None])
+        layer_2.append(code_list[7 * i + 1][None] - 4096)
+        layer_3.append(code_list[7 * i + 2][None] - (2 * 4096))
+        layer_3.append(code_list[7 * i + 3][None] - (3 * 4096))
+        layer_2.append(code_list[7 * i + 4][None] - (4 * 4096))
+        layer_3.append(code_list[7 * i + 5][None] - (5 * 4096))
+        layer_3.append(code_list[7 * i + 6][None] - (6 * 4096))
+    
+    with torch.no_grad():
+        codes = [
+            torch.concat(layer_1),
+            torch.concat(layer_2),
+            torch.concat(layer_3)
+        ]
+        for i in range(len(codes)):
+            codes[i][codes[i] < 0] = 0
+            codes[i] = codes[i][None]
+        
+        audio_hat = snac_model.decode(codes)
+        return audio_hat.cpu()[0, 0]
 
-sg2_ckpt = "google/siglip2-so400m-patch14-384"
-siglip2_model = AutoModel.from_pretrained(sg2_ckpt, device_map="cpu").eval()
-siglip2_processor = AutoProcessor.from_pretrained(sg2_ckpt)
+# Load the SNAC model (shared by all)
+snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cuda")
 
-@spaces.GPU
-def postprocess_siglip(sg1_probs, sg2_probs, labels):
-    sg1_output = {labels[i]: sg1_probs[0][i].item() for i in range(len(labels))}
-    sg2_output = {labels[i]: sg2_probs[0][i].item() for i in range(len(labels))}
-    return sg1_output, sg2_output
+# Load all the single-speaker language models
+models = {
+    "Luna": {
+        "tokenizer": AutoTokenizer.from_pretrained('prithivMLmods/Llama-3B-Mono-Luna'),
+        "model": AutoModelForCausalLM.from_pretrained('prithivMLmods/Llama-3B-Mono-Luna', torch_dtype=torch.bfloat16).cuda()
+    },
+    "Ceylia": {
+        "tokenizer": AutoTokenizer.from_pretrained('prithivMLmods/Llama-3B-Mono-Ceylia'),
+        "model": AutoModelForCausalLM.from_pretrained('prithivMLmods/Llama-3B-Mono-Ceylia', torch_dtype=torch.bfloat16).cuda()
+    },
+    "Cooper": {
+        "tokenizer": AutoTokenizer.from_pretrained('prithivMLmods/Llama-3B-Mono-Cooper'),
+        "model": AutoModelForCausalLM.from_pretrained('prithivMLmods/Llama-3B-Mono-Cooper', torch_dtype=torch.bfloat16).cuda()
+    },
+    "Jim": {
+        "tokenizer": AutoTokenizer.from_pretrained('prithivMLmods/Llama-3B-Mono-Jim'),
+        "model": AutoModelForCausalLM.from_pretrained('prithivMLmods/Llama-3B-Mono-Jim', torch_dtype=torch.bfloat16).cuda()
+    },
+}
 
-def siglip_detector(image, texts):
-    sg1_inputs = siglip1_processor(
-        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
-    ).to("cpu")
-    sg2_inputs = siglip2_processor(
-        text=texts, images=image, return_tensors="pt", padding="max_length", max_length=64
-    ).to("cpu")
+def generate_audio(text, temperature, top_p, max_new_tokens, model_name):
+    """
+    Given input text and model parameters, generate speech audio using the chosen model.
+    """
+    # Retrieve the chosen tokenizer and model
+    chosen = models[model_name]
+    tokenizer = chosen["tokenizer"]
+    model = chosen["model"]
+    
+    prompt = f'<custom_token_3><|begin_of_text|>{text}<|eot_id|><custom_token_4><custom_token_5><custom_token_1>'
+    input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt').to('cuda')
+    
     with torch.no_grad():
-        sg1_outputs = siglip1_model(**sg1_inputs)
-        sg2_outputs = siglip2_model(**sg2_inputs)
-        sg1_logits_per_image = sg1_outputs.logits_per_image
-        sg2_logits_per_image = sg2_outputs.logits_per_image
-        sg1_probs = torch.sigmoid(sg1_logits_per_image)
-        sg2_probs = torch.sigmoid(sg2_logits_per_image)
-    return sg1_probs, sg2_probs
-
-def infer(image, candidate_labels):
-    candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")]
-    sg1_probs, sg2_probs = siglip_detector(image, candidate_labels)
-    return postprocess_siglip(sg1_probs, sg2_probs, labels=candidate_labels)
-
-# Build the Gradio Interface with two tabs.
-with gr.Blocks(theme="YTheme/Minecraft") as demo:
-    gr.Markdown("# Multi-Domain & Zero-Shot Image Classification")
+        generated_ids = model.generate(
+            **input_ids,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=1.1,
+            num_return_sequences=1,
+            eos_token_id=128258,
+        )
     
-    with gr.Tabs():
-        # Tab 1: Multi-Model Classification
-        with gr.Tab("Multi-Domain Classification"):
-            with gr.Sidebar():
-                gr.Markdown("# Choose Domain")
-                with gr.Row():
-                    age_btn = gr.Button("Age Classification", variant="primary")
-                    gender_btn = gr.Button("Gender Classification", variant="secondary")
-                    emotion_btn = gr.Button("Emotion Classification", variant="secondary")
-                    gym_workout_btn = gr.Button("Gym Workout Classification", variant="secondary")
-                    dog_breed_btn = gr.Button("Dog Breed Classification", variant="secondary")
-                    bird_btn = gr.Button("Bird Species Classification", variant="secondary")
-                    waste_btn = gr.Button("Waste Classification", variant="secondary")
-                    deepfake_btn = gr.Button("Deepfake Quality Test", variant="secondary")
-                    traffic_density_btn = gr.Button("Traffic Density", variant="secondary")
-                    sign_language_btn = gr.Button("Alphabet Sign Language", variant="secondary")
-                    clip_art_btn = gr.Button("Clip Art 126", variant="secondary")
-                    mnist_btn = gr.Button("Digit Classify (0-9)", variant="secondary")
-                    fashion_mnist_btn = gr.Button("Fashion MNIST (only cloth)", variant="secondary")
-                    food_btn = gr.Button("Indian/Western Food Type", variant="secondary")
-                    leaf_disease_btn = gr.Button("Rice Leaf Disease", variant="secondary")
-                    multisource_btn = gr.Button("Multi Source 121", variant="secondary")
-                    painting_btn = gr.Button("Painting 126", variant="secondary")
-                    sketch_btn = gr.Button("Sketch 126", variant="secondary")
-                    
-                selected_model = gr.State("age")
-                gr.Markdown("### Current Model:")
-                model_display = gr.Textbox(value="age", interactive=False)
-                selected_model.change(lambda m: m, selected_model, model_display)
+    row = generated_ids[0, input_ids['input_ids'].shape[1]:]
+    y_tensor = redistribute_codes(row)
+    y_np = y_tensor.detach().cpu().numpy()
+    return (24000, y_np)
 
-                buttons = [
-                    gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn,
-                    age_btn, mnist_btn, fashion_mnist_btn, food_btn, bird_btn, leaf_disease_btn,
-                    sign_language_btn, traffic_density_btn, clip_art_btn, multisource_btn, painting_btn, sketch_btn  # Include new button
-                ]
-                model_names = [
-                    "gender", "emotion", "dog breed", "deepfake", "gym workout", "waste",
-                    "age", "mnist", "fashion_mnist", "food", "bird", "leaf disease",
-                    "sign language", "traffic density", "clip art", "multisource", "painting", "sketch"  # New model name
-                ]
+# Example texts with emotion tokens
+example_texts = [
+    ["Hi, my name is Alex. <laugh> It's a wonderful day! <chuckle> I love coding."],
+    ["I woke up feeling sleepy. <yawn> I need coffee! <sniffle> But I'm ready to work."],
+    ["Oh no, I forgot my keys! <groan> <uhm> Maybe I'll try again later. <sigh>"],
+    ["This is amazing! <gasp> Really, it's fantastic. <giggles>"]
+]
 
-                for btn, name in zip(buttons, model_names):
-                    btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
-
-            with gr.Row():
-                with gr.Column():
-                    image_input = gr.Image(type="numpy", label="Upload Image")
-                    analyze_btn = gr.Button("Classify / Predict")
-                output_label = gr.Label(label="Prediction Scores")
-                analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)
-        
-        # Tab 2: Zero-Shot Classification (SigLIP)
-        with gr.Tab("Zero-Shot Classification"):
-            gr.Markdown("## Compare SigLIP 1 and SigLIP 2 on Zero-Shot Classification")
-            with gr.Row():
-                with gr.Column():
-                    zs_image_input = gr.Image(type="pil", label="Upload Image")
-                    zs_text_input = gr.Textbox(label="Input a list of labels (comma separated)")
-                    zs_run_button = gr.Button("Run")
-                with gr.Column():
-                    siglip1_output = gr.Label(label="SigLIP 1 Output", num_top_classes=3)
-                    siglip2_output = gr.Label(label="SigLIP 2 Output", num_top_classes=3)
-            zs_run_button.click(fn=infer, inputs=[zs_image_input, zs_text_input], outputs=[siglip1_output, siglip2_output])
+# Gradio Interface
+with gr.Blocks() as demo:
+    # Sidebar for model selection
+    with gr.Sidebar():
+        gr.Markdown("# Choose Model")
+        model_choice = gr.Dropdown(choices=list(models.keys()), value="Luna", label="Model")
+    
+    gr.Markdown("# Single Speaker Audio Generation")
+    gr.Markdown("Generate speech audio using one of the single-speaker models. Use the examples below to see how emotion tokens like `<laugh>`, `<chuckle>`, `<sigh>`, etc. can be incorporated.")
+    
+    with gr.Row():
+        text_input = gr.Textbox(lines=4, label="Input Text")
+    
+    # Examples with emotion tokens
+    gr.Examples(
+        examples=example_texts,
+        inputs=text_input,
+        label="Emotion Examples",
+        cache_examples=False
+    )
+    
+    with gr.Row():
+        temp_slider = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=0.9, label="Temperature")
+        top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.8, label="Top-p")
+        tokens_slider = gr.Slider(minimum=100, maximum=3500, step=50, value=1200, label="Max New Tokens")
+    
+    output_audio = gr.Audio(type="numpy", label="Generated Audio")
+    generate_button = gr.Button("Generate Audio")
+    
+    # Pass the selected model name along with other parameters
+    generate_button.click(
+        fn=generate_audio,
+        inputs=[text_input, temp_slider, top_p_slider, tokens_slider, model_choice],
+        outputs=output_audio
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()