# ============================================================
# ControlFoley - Hugging Face ZeroGPU Gradio App
# Python 3.10
# PyTorch 2.8 / ZeroGPU compatible
# ============================================================

# IMPORTANT:
# Hugging Face ZeroGPU requires importing spaces BEFORE torch.
import spaces

import os
import sys
import time
import shutil
import logging
import subprocess
import tempfile
from pathlib import Path
from contextlib import contextmanager

import gradio as gr
from huggingface_hub import snapshot_download

import torch
import torchaudio


# ============================================================
# Paths / configuration
# ============================================================

APP_DIR = Path(__file__).resolve().parent

SOURCE_DIR = APP_DIR / "upstream_controlfoley"
MODEL_DIR = APP_DIR / "model_weights"
OUTPUT_DIR = APP_DIR / "outputs"
ASSET_DIR = APP_DIR / "sample_assets"

UPSTREAM_REPO = "https://github.com/xiaomi-research/controlfoley.git"
MODEL_REPO = "YJX-Xiaomi/ControlFoley"

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
ASSET_DIR.mkdir(parents=True, exist_ok=True)

# ControlFoley contains several relative paths internally.
# Make sure its expected working directory is the Space root.
os.chdir(APP_DIR)


# ============================================================
# Download official ControlFoley source
# ============================================================

def clone_upstream():
    if (SOURCE_DIR / "controlfoley").exists():
        print("ControlFoley source already available.")
        return

    print("Cloning official ControlFoley repository...")

    subprocess.run(
        [
            "git",
            "clone",
            "--depth",
            "1",
            UPSTREAM_REPO,
            str(SOURCE_DIR),
        ],
        check=True,
    )


# ============================================================
# Download official ControlFoley model weights
# ============================================================

def download_model_files():
    required_main = MODEL_DIR / "weights" / "controlfoley.pth"

    required_ext = [
        MODEL_DIR / "ext_weights" / "v1-44.pth",
        MODEL_DIR / "ext_weights" / "synchformer_state_dict.pth",
        MODEL_DIR / "ext_weights" / "cav_mae_st.pth",
        MODEL_DIR
        / "ext_weights"
        / "music_speech_audioset_epoch_15_esc_89.98.pt",
    ]

    if required_main.exists() and all(x.exists() for x in required_ext):
        print("ControlFoley model files already available.")
        return

    print("Downloading ControlFoley model weights...")

    snapshot_download(
        repo_id=MODEL_REPO,
        local_dir=str(MODEL_DIR),
        allow_patterns=[
            "weights/*",
            "ext_weights/*",
        ],
    )


# ============================================================
# Startup downloads
# ============================================================

clone_upstream()
download_model_files()


# ============================================================
# Configure Python paths
# ============================================================

# ControlFoley imports modules from both repository root and lib/.
if str(SOURCE_DIR) not in sys.path:
    sys.path.insert(0, str(SOURCE_DIR))

if str(SOURCE_DIR / "lib") not in sys.path:
    sys.path.insert(0, str(SOURCE_DIR / "lib"))


# ============================================================
# Import ControlFoley
# ============================================================

from controlfoley.inference_utils import (
    all_model_cfg,
    generate,
    load_video,
    make_video,
    setup_eval_logging,
)

from controlfoley.audio_model import (
    create_audio_generation_model,
)

from controlfoley.feature_extractor import FeaturesUtils

from lib.flow_matching import FlowMatching


# ============================================================
# Logging
# ============================================================

setup_eval_logging()

log = logging.getLogger("controlfoley-space")

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


# ============================================================
# PyTorch 2.6+ compatibility
# ============================================================
#
# PyTorch >=2.6 changed:
#
#     torch.load(..., weights_only=False)
#
# from effectively the old default behavior to:
#
#     weights_only=True
#
# Some older libraries bundled/used by ControlFoley, especially
# LAION-CLAP and AudioCraft/MusicGen, call torch.load() without
# specifying weights_only.
#
# Their official checkpoints contain more than plain tensor
# state dictionaries, so weights_only=True fails.
#
# We DO NOT monkey-patch torch.load globally.
#
# Instead, this compatibility context is active ONLY while
# initializing trusted upstream ControlFoley feature models.
# ============================================================

@contextmanager
def legacy_checkpoint_loading():
    original_torch_load = torch.load

    def compatible_torch_load(*args, **kwargs):
        # Preserve an explicit value supplied by a library.
        #
        # Only restore the old behavior when the caller does
        # not specify weights_only at all.
        if "weights_only" not in kwargs:
            kwargs["weights_only"] = False

        return original_torch_load(*args, **kwargs)

    torch.load = compatible_torch_load

    try:
        yield
    finally:
        torch.load = original_torch_load


# ============================================================
# ControlFoley configuration
# ============================================================

MODEL_CFG = all_model_cfg["large_44k"]
SEQ_CFG = MODEL_CFG.seq_cfg


# ============================================================
# Main ControlFoley network
# ============================================================

print("=" * 60)
print("Loading ControlFoley main network...")
print("=" * 60)

NET = create_audio_generation_model(
    MODEL_CFG.model_name
).to(
    "cuda",
    torch.float32,
).eval()


MAIN_CHECKPOINT = MODEL_DIR / "weights" / "controlfoley.pth"

# The main ControlFoley checkpoint is a normal weight state dict,
# therefore keeping weights_only=True is appropriate here.
main_state = torch.load(
    MAIN_CHECKPOINT,
    map_location="cuda",
    weights_only=True,
)

NET.load_weights(main_state)

del main_state

print("ControlFoley main network loaded.")


# ============================================================
# Feature extractor stack
# ============================================================
#
# Includes:
#
# - DFN5B CLIP
# - Synchformer
# - CAV-MAE-ST
# - LAION CLAP
# - AudioCraft MusicGen Style
# - VAE / vocoder
#
# LAION CLAP and some AudioCraft checkpoints require the legacy
# PyTorch checkpoint-loading behavior.
# ============================================================

print("=" * 60)
print("Loading ControlFoley feature extractors...")
print("=" * 60)

with legacy_checkpoint_loading():

    FEATURES = FeaturesUtils(
        tod_vae_ckpt=str(
            MODEL_DIR
            / "ext_weights"
            / "v1-44.pth"
        ),

        synchformer_ckpt=str(
            MODEL_DIR
            / "ext_weights"
            / "synchformer_state_dict.pth"
        ),

        cav_mae_ckpt=str(
            MODEL_DIR
            / "ext_weights"
            / "cav_mae_st.pth"
        ),

        clap_ckpt=str(
            MODEL_DIR
            / "ext_weights"
            / "music_speech_audioset_epoch_15_esc_89.98.pt"
        ),

        mode=MODEL_CFG.mode,
        enable_conditions=True,
        need_vae_encoder=False,
    )


FEATURES = FEATURES.to(
    "cuda",
    torch.float32,
).eval()

print("ControlFoley feature extractors loaded.")


# ============================================================
# Sample video
# ============================================================

def prepare_sample_video():
    target = ASSET_DIR / "001.mp4"

    if target.exists():
        return target

    upstream_sample = SOURCE_DIR / "assets" / "001.mp4"

    if not upstream_sample.exists():
        print("Sample video not found.")
        return None

    shutil.copy2(
        upstream_sample,
        target,
    )

    return target


SAMPLE_VIDEO = prepare_sample_video()


# ============================================================
# Dynamic ZeroGPU duration
# ============================================================

def gpu_budget(
    video_path,
    prompt,
    negative_prompt,
    duration,
    cfg_strength,
    steps,
    seed,
):
    """
    Allocate enough ZeroGPU time depending on inference settings.
    """

    duration = float(duration)
    steps = int(steps)

    estimated = (
        60
        + int(duration * 12)
        + int(steps * 3)
    )

    # ZeroGPU maximum requested allocation.
    return min(
        300,
        max(120, estimated),
    )


# ============================================================
# Main generation function
# ============================================================

@spaces.GPU(
    size="xlarge",
    duration=gpu_budget,
)
@torch.inference_mode()
def generate_foley(
    video_path,
    prompt,
    negative_prompt,
    duration,
    cfg_strength,
    steps,
    seed,
):

    if not video_path:
        raise gr.Error(
            "Please upload a video or select the sample video."
        )

    # --------------------------------------------------------
    # Parameters
    # --------------------------------------------------------

    video_path = Path(video_path)

    duration = float(duration)
    cfg_strength = float(cfg_strength)
    steps = int(steps)
    seed = int(seed)

    if duration < 1 or duration > 8:
        raise gr.Error(
            "Duration must be between 1 and 8 seconds."
        )

    if steps < 5 or steps > 30:
        raise gr.Error(
            "Inference steps must be between 5 and 30."
        )

    if not video_path.exists():
        raise gr.Error(
            "The uploaded video could not be found."
        )


    # --------------------------------------------------------
    # Output directory
    # --------------------------------------------------------

    job_dir = Path(
        tempfile.mkdtemp(
            prefix="controlfoley_",
            dir=str(OUTPUT_DIR),
        )
    )

    audio_path = (
        job_dir
        / "generated_foley.flac"
    )

    video_out_path = (
        job_dir
        / "video_with_generated_audio.mp4"
    )


    # --------------------------------------------------------
    # Load / preprocess video
    # --------------------------------------------------------

    print(f"Loading video: {video_path}")

    video_info = load_video(
        video_path,
        duration,
    )

    actual_duration = min(
        duration,
        float(video_info.total_duration),
    )


    # --------------------------------------------------------
    # Video conditioning
    # --------------------------------------------------------

    clip_frames = (
        video_info
        .clip_embeddings
        .unsqueeze(0)
    )

    visual_frames = (
        video_info
        .visual_features
        .unsqueeze(0)
    )

    sync_frames = (
        video_info
        .sync_embeddings
        .unsqueeze(0)
    )


    # --------------------------------------------------------
    # Configure temporal dimensions
    # --------------------------------------------------------

    SEQ_CFG.total_time_seconds = actual_duration

    NET.update_seq_lengths(
        SEQ_CFG.latent_sequence_length,
        SEQ_CFG.clip_sequence_length,
        SEQ_CFG.visual_sequence_length,
        SEQ_CFG.sync_sequence_length,
    )


    # --------------------------------------------------------
    # Random generator
    # --------------------------------------------------------

    rng = torch.Generator(
        device="cuda"
    )

    rng.manual_seed(seed)


    # --------------------------------------------------------
    # Flow matching sampler
    # --------------------------------------------------------

    fm = FlowMatching(
        min_sigma=0,
        inference_mode="euler",
        num_steps=steps,
    )


    # --------------------------------------------------------
    # Generate
    # --------------------------------------------------------

    print("=" * 60)
    print("Generating Foley audio...")
    print(f"Prompt: {prompt}")
    print(f"Duration: {actual_duration}")
    print(f"Steps: {steps}")
    print(f"CFG: {cfg_strength}")
    print(f"Seed: {seed}")
    print("=" * 60)

    start_time = time.time()


    audios = generate(

        # Video conditioning
        clip_frames,
        visual_frames,
        sync_frames,

        # No reference audio
        None,

        # No timbre reference
        None,

        # Reference audio duration
        0.0,

        # Text prompt
        [prompt or ""],

        negative_text=[
            negative_prompt or ""
        ],

        feature_utils=FEATURES,
        net=NET,
        fm=fm,
        rng=rng,
        cfg_strength=cfg_strength,
    )


    # --------------------------------------------------------
    # Convert generated tensor
    # --------------------------------------------------------

    audio = (
        audios
        .float()
        .cpu()[0]
    )


    # --------------------------------------------------------
    # Save generated FLAC
    # --------------------------------------------------------

    torchaudio.save(
        str(audio_path),
        audio,
        SEQ_CFG.audio_sample_rate,
    )


    # --------------------------------------------------------
    # Mux generated audio with original video
    # --------------------------------------------------------

    make_video(
        video_info,
        video_out_path,
        audio,
        sampling_rate=SEQ_CFG.audio_sample_rate,
    )


    # --------------------------------------------------------
    # Finished
    # --------------------------------------------------------

    elapsed = time.time() - start_time

    status = f"""
### ✅ Generation complete

**Duration:** {actual_duration:.2f}s  
**Generation time:** {elapsed:.1f}s  
**Steps:** {steps}  
**CFG:** {cfg_strength}  
**Seed:** {seed}
"""


    # --------------------------------------------------------
    # Release temporary tensors
    # --------------------------------------------------------

    del audios
    del audio

    try:
        del clip_frames
        del visual_frames
        del sync_frames
    except Exception:
        pass

    if torch.cuda.is_available():
        torch.cuda.empty_cache()


    return (
        str(audio_path),
        str(video_out_path),
        status,
    )


# ============================================================
# Gradio UI
# ============================================================

TITLE = """
# 🎬 ControlFoley — Video → Foley Audio

Generate synchronized Foley sound effects from video using
**Xiaomi Research ControlFoley**.

Upload a video or select the official sample below.

You can use:

- **V2A** — leave the prompt empty
- **TV2A** — describe the sound you want
"""


HELP_TEXT = """
### Usage

**Pure Video → Audio**

Leave the prompt blank.

**Text-guided Video → Audio**

Describe the expected sound.

Example:

`the skateboard wheels scraping and grinding on the ground.`

For the first test, use:

- Duration: **8 seconds**
- Steps: **25**
- CFG: **4.5**
- Seed: **42**

If you want a quicker test, reduce inference steps to **10–15**.
"""


with gr.Blocks(
    title="ControlFoley Video to Audio"
) as demo:

    gr.Markdown(TITLE)

    with gr.Row():

        # ====================================================
        # INPUT
        # ====================================================

        with gr.Column():

            video = gr.Video(
                label="Input Video",
                sources=["upload"],
                format="mp4",
            )

            prompt = gr.Textbox(
                label="Sound Prompt",
                placeholder=(
                    "Describe the sound, or leave blank "
                    "for pure Video-to-Audio"
                ),
                value=(
                    "the skateboard wheels scraping "
                    "and grinding on the ground."
                ),
                lines=2,
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder=(
                    "Example: music, speech, crowd noise"
                ),
                value="",
                lines=1,
            )


            with gr.Accordion(
                "Generation Settings",
                open=False,
            ):

                duration = gr.Slider(
                    minimum=1,
                    maximum=8,
                    value=8,
                    step=0.5,
                    label="Duration (seconds)",
                )

                cfg_strength = gr.Slider(
                    minimum=1.0,
                    maximum=8.0,
                    value=4.5,
                    step=0.5,
                    label="CFG Strength",
                )

                steps = gr.Slider(
                    minimum=5,
                    maximum=30,
                    value=25,
                    step=1,
                    label="Inference Steps",
                )

                seed = gr.Number(
                    value=42,
                    precision=0,
                    label="Seed",
                )


            generate_btn = gr.Button(
                "Generate Foley Sound",
                variant="primary",
            )


        # ====================================================
        # OUTPUT
        # ====================================================

        with gr.Column():

            audio_out = gr.Audio(
                label="Generated Foley Audio",
                type="filepath",
            )

            video_out = gr.Video(
                label="Video + Generated Foley",
            )

            status = gr.Markdown()


    # ========================================================
    # Sample
    # ========================================================

    if SAMPLE_VIDEO is not None:

        gr.Examples(
            examples=[
                [
                    str(SAMPLE_VIDEO),
                    (
                        "the skateboard wheels scraping "
                        "and grinding on the ground."
                    ),
                    "",
                    8,
                    4.5,
                    25,
                    42,
                ],
            ],
            inputs=[
                video,
                prompt,
                negative_prompt,
                duration,
                cfg_strength,
                steps,
                seed,
            ],
            label="Official ControlFoley Sample",
        )


    gr.Markdown(HELP_TEXT)


    # ========================================================
    # Generation event
    # ========================================================

    generate_btn.click(
        fn=generate_foley,

        inputs=[
            video,
            prompt,
            negative_prompt,
            duration,
            cfg_strength,
            steps,
            seed,
        ],

        outputs=[
            audio_out,
            video_out,
            status,
        ],

        api_name="generate_foley",

        concurrency_limit=1,
    )


# ============================================================
# Launch
# ============================================================

if __name__ == "__main__":

    demo.queue(
        max_size=20,
        default_concurrency_limit=1,
    ).launch()