import spaces
import gradio as gr
import torch
import argparse
from seed_vc_wrapper import SeedVCWrapper

# Set up device and torch configurations
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.triton.unique_kernel_names = True

if hasattr(torch._inductor.config, "fx_graph_cache"):
    # Experimental feature to reduce compilation times, will be on by default in future
    torch._inductor.config.fx_graph_cache = True

dtype = torch.float16

# Global variables to store model instances
vc_wrapper_v1 = SeedVCWrapper()

@spaces.GPU
def convert_voice_v1_wrapper(source_audio_path, target_audio_path, diffusion_steps=10,
                             length_adjust=1.0, inference_cfg_rate=0.7,
                             auto_f0_adjust=True, pitch_shift=0, stream_output=True):
    """
    Wrapper function for vc_wrapper.convert_voice that can be decorated with @spaces.GPU
    """

    # Use yield from to properly handle the generator
    yield from vc_wrapper_v1.convert_voice(
        source=source_audio_path,
        target=target_audio_path,
        diffusion_steps=diffusion_steps,
        length_adjust=length_adjust,
        inference_cfg_rate=inference_cfg_rate,
        f0_condition=True,  # Always True as requested - removed from UI
        auto_f0_adjust=auto_f0_adjust,
        pitch_shift=pitch_shift,
        stream_output=stream_output
    )

def create_v1_interface():
    # Set up Gradio interface
    description = (
        "<b>Zero shot voice conversion across all Indian languages</b>, achieved by finetuning a Seed-VoiceConversion checkpoint with Indic datasets. <br> "
        "For instructions on <b>local deployment</b> and further finetuning, please refer [<b>Plachtaa/seed-vc</b>](https://github.com/Plachtaa/seed-vc) . The finetuned checkpoints are available for download on our [<b>model page</b>](https://huggingface.co/DreamSyncCo/IndicVoiceChanger). <br>"
        "<b>Note:</b> Any reference audio will be forcefully clipped to <b>25s</b> if beyond this length.<br> "
        "If total duration of source and reference audio exceeds <b>30s</b>, source audio will be processed in chunks.<br>")

    inputs = [
        gr.Audio(type="filepath", label="Source Audio"),
        gr.Audio(type="filepath", label="Reference Audio"),
        gr.Slider(minimum=1, maximum=200, value=10, step=1, label="Diffusion Steps",
                  info="10 by default, 50~100 for best quality"),
        gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Adjust",
                  info="<1.0 for speed-up speech, >1.0 for slow-down speech"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Inference CFG Rate",
                  info="has subtle influence"),
        gr.Checkbox(label="Auto F0 adjust", value=True,
                    info="Roughly adjust F0 to match target voice."),
        gr.Slider(label='Pitch shift', minimum=-24, maximum=24, step=1, value=0,
                  info="Pitch shift in semitones, only works when F0 conditioned model is used"),
    ]

    examples = [
        ["examples/source/Hindi.wav", "examples/reference/Marathi.wav", 25, 1.0, 0.7, True, 0],
        ["examples/source/Assamese.wav", "examples/reference/Kannada.wav", 25, 1.0, 0.7, False, 0],
        ["examples/source/Malayalam.wav", "examples/reference/Telugu.wav", 25, 1.0, 0.7, False, 0],
        ["examples/source/Tamil.wav", "examples/reference/Bengali.wav", 25, 1.0, 0.7, True, 0],
        
        
    ]

    outputs = [
        gr.Audio(label="Stream Output Audio", streaming=True, format='mp3'),
        gr.Audio(label="Full Output Audio", streaming=False, format='wav')
    ]

    return gr.Interface(
        fn=convert_voice_v1_wrapper,
        description=description,
        inputs=inputs,
        outputs=outputs,
        title="<b>Voice Conversion for Indian Languages</b>",
        examples=examples,
        cache_examples=False,
    )


def main(args):
    # Create interface
    v1_interface = create_v1_interface()

    # Launch the interface
    v1_interface.launch()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--compile", type=bool, default=True)
    args = parser.parse_args()
    main(args)