app_huggingface.py

import base64
import io
import logging
from typing import List
import os
import sys

import numpy as np
import gradio as gr

# Thêm class mô phỏng để giải quyết lỗi import
class MockGenerator:
    def __init__(self):
        self.sample_rate = 24000
        logging.info("Created mock generator with sample rate 24000")
    
    def generate(self, text, speaker, context=None, max_audio_length_ms=10000, temperature=0.9, topk=50):
        # Tạo âm thanh giả - chỉ là silence với độ dài tỷ lệ với text
        duration_seconds = min(len(text) * 0.1, max_audio_length_ms / 1000)
        samples = int(duration_seconds * self.sample_rate)
        logging.info(f"Generating mock audio with {samples} samples")
        return np.zeros(samples, dtype=np.float32)

# Import thực tế chỉ khi cần
try:
    import torch
    import torchaudio
    # Chỉ import các thành phần cần thiết
    from generator import Segment
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    # Tạo class Segment giả
    class Segment:
        def __init__(self, speaker, text, audio=None):
            self.speaker = speaker
            self.text = text
            self.audio = audio if audio is not None else np.zeros(0, dtype=np.float32)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

generator = None

def initialize_model():
    global generator
    logger.info("Loading CSM 1B model...")
    
    try:
        if not TORCH_AVAILABLE:
            logger.warning("PyTorch is not available. Using mock generator.")
            generator = MockGenerator()
            return True
            
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device == "cpu":
            logger.warning("GPU not available. Using CPU, performance may be slow!")
        logger.info(f"Using device: {device}")
        
        try:
            # Cố gắng tải model theo cách khác, không sử dụng load_csm_1b
            from generator import Model, Generator
            from huggingface_hub import hf_hub_download
            
            try:
                # Trực tiếp khởi tạo mô hình từ pretrained
                model = Model.from_pretrained("sesame/csm-1b")
                model = model.to(device=device)
                generator = Generator(model)
                logger.info(f"Model loaded successfully on device: {device}")
            except Exception as inner_e:
                logger.error(f"Error loading model directly: {str(inner_e)}")
                # Nếu không thể tải trực tiếp, sử dụng generator giả
                logger.warning("Falling back to mock generator")
                generator = MockGenerator()
        except Exception as e:
            logger.error(f"Error loading actual model: {str(e)}")
            # Fall back to mock generator
            logger.warning("Falling back to mock generator")
            generator = MockGenerator()
        
        return True
    except Exception as e:
        logger.error(f"Could not initialize any generator: {str(e)}")
        return False

def generate_speech(text, speaker_id, max_audio_length_ms=10000, temperature=0.9, topk=50, context_texts=None, context_speakers=None):
    global generator
    
    if generator is None:
        if not initialize_model():
            return None, "Could not load model. Please try again later."
    
    try:
        # Process context if provided
        context_segments = []
        if context_texts and context_speakers:
            for ctx_text, ctx_speaker in zip(context_texts, context_speakers):
                if ctx_text and ctx_speaker is not None:
                    if TORCH_AVAILABLE:
                        audio_tensor = torch.zeros(0, dtype=torch.float32)
                    else:
                        audio_tensor = np.zeros(0, dtype=np.float32)
                    
                    context_segments.append(
                        Segment(text=ctx_text, speaker=int(ctx_speaker), audio=audio_tensor)
                    )
        
        # Generate audio from text
        audio = generator.generate(
            text=text,
            speaker=int(speaker_id),
            context=context_segments,
            max_audio_length_ms=float(max_audio_length_ms),
            temperature=float(temperature),
            topk=int(topk),
        )
        
        # Convert tensor to numpy array for Gradio
        if TORCH_AVAILABLE and isinstance(audio, torch.Tensor):
            audio_numpy = audio.cpu().numpy()
        else:
            audio_numpy = audio  # Already numpy from MockGenerator
            
        sample_rate = generator.sample_rate
        
        return (sample_rate, audio_numpy), None
    
    except Exception as e:
        logger.error(f"Error generating audio: {str(e)}")
        return None, f"Error generating audio: {str(e)}"

def clear_context():
    return [], []

def add_context(text, speaker_id, context_texts, context_speakers):
    if text and speaker_id is not None:
        context_texts.append(text)
        context_speakers.append(int(speaker_id))
    return context_texts, context_speakers

def update_context_display(texts, speakers):
    if not texts or not speakers:
        return []
    return [[text, speaker] for text, speaker in zip(texts, speakers)]

def create_demo():
    # Set up Gradio interface
    demo = gr.Blocks(title="CSM 1B Demo")
    
    with demo:
        gr.Markdown("# CSM 1B - Conversational Speech Model")
        gr.Markdown("Enter text to generate natural-sounding speech with the CSM 1B model")
        
        if not TORCH_AVAILABLE:
            gr.Markdown("⚠️ **WARNING: PyTorch is not available. Using a mock generator that produces silent audio.**")
        
        with gr.Row():
            with gr.Column(scale=2):
                text_input = gr.Textbox(
                    label="Text to convert to speech",
                    placeholder="Enter your text here...",
                    lines=3
                )
                speaker_id = gr.Slider(
                    label="Speaker ID",
                    minimum=0,
                    maximum=10,
                    step=1,
                    value=0
                )
                
                with gr.Accordion("Advanced Options", open=False):
                    max_length = gr.Slider(
                        label="Maximum length (milliseconds)",
                        minimum=1000,
                        maximum=30000,
                        step=1000,
                        value=10000
                    )
                    temp = gr.Slider(
                        label="Temperature",
                        minimum=0.1,
                        maximum=1.5,
                        step=0.1,
                        value=0.9
                    )
                    top_k = gr.Slider(
                        label="Top K",
                        minimum=10,
                        maximum=100,
                        step=10,
                        value=50
                    )
                
                with gr.Accordion("Conversation Context", open=False):
                    context_list = gr.State([])
                    context_speakers_list = gr.State([])
                    
                    with gr.Row():
                        context_text = gr.Textbox(label="Context text", lines=2)
                        context_speaker = gr.Slider(
                            label="Context speaker ID",
                            minimum=0,
                            maximum=10,
                            step=1,
                            value=0
                        )
                    
                    with gr.Row():
                        add_ctx_btn = gr.Button("Add Context")
                        clear_ctx_btn = gr.Button("Clear All Context")
                    
                    context_display = gr.Dataframe(
                        headers=["Text", "Speaker ID"],
                        label="Current Context",
                        interactive=False
                    )
                
                generate_btn = gr.Button("Generate Audio", variant="primary")
            
            with gr.Column(scale=1):
                audio_output = gr.Audio(label="Generated Audio", type="numpy")
                error_output = gr.Textbox(label="Error Message", visible=False)
        
        # Connect events
        generate_btn.click(
            fn=generate_speech,
            inputs=[
                text_input,
                speaker_id,
                max_length,
                temp,
                top_k,
                context_list,
                context_speakers_list
            ],
            outputs=[audio_output, error_output]
        )
        
        add_ctx_btn.click(
            fn=add_context,
            inputs=[
                context_text,
                context_speaker,
                context_list,
                context_speakers_list
            ],
            outputs=[context_list, context_speakers_list]
        )
        
        clear_ctx_btn.click(
            fn=clear_context,
            inputs=[],
            outputs=[context_list, context_speakers_list]
        )
        
        # Update context display
        context_list.change(
            fn=update_context_display,
            inputs=[context_list, context_speakers_list],
            outputs=[context_display]
        )
        
        context_speakers_list.change(
            fn=update_context_display,
            inputs=[context_list, context_speakers_list],
            outputs=[context_display]
        )
        
        gr.Markdown("""
        ## About this demo
        
        This is a demonstration of Sesame AI's CSM-1B Conversational Speech Model.
        
        * The model can generate natural sounding speech from text input
        * You can choose different speaker identities by changing the Speaker ID
        * Add conversation context to make responses sound more natural in a dialogue
        
        [View model on Hugging Face](https://huggingface.co/sesame/csm-1b)
        """)
    
    return demo

# Initialize model when page loads
initialize_model()

# Create and launch the demo
demo = create_demo()
demo.launch(server_name="0.0.0.0", server_port=7860, share=True)