"""Graphics Studio on ZeroGPU.

Image tab (text to image, image generation): FLUX.2 klein 9B turns a prompt into a still image.
Music tab (text to music, text to audio): MusicGen medium turns a description into up to 28 seconds of music.
Transcribe tab (speech to text, speech recognition, automatic speech recognition): Whisper large v3 turbo
writes down speech from an uploaded file or the microphone, in its own language or translated to English.
Narration tab (text to speech, machine speech, speech generation): Kokoro reads your text aloud.
Describe tab: three sub-tabs for the describing jobs the Space does.
  Image to text (image captioning, image recognition): Florence-2 large writes a detailed caption for
    an uploaded picture.
  Video to text (video captioning, video to text, video-text-to-text): SmolVLM2 2.2B looks at frames
    from an uploaded clip and describes the scene and the motion.
  Music to text (music captioning, audio captioning, music classification, audio classification):
    CLAP picks a genre, mood, tempo and instruments for a piece of music and writes them as one short
    sentence.
Video tab (text to video, image to video, video generation): one button makes a whole clip. A small
    text model (Qwen3 4B) first improves the scene and motion prompts if that is ticked, FLUX draws
    the scene, LTX-Video (13B) animates it, MusicGen medium adds a soundtrack if you describe one,
    and if narration is ticked Qwen3 4B writes the lines and their timing and Kokoro speaks them.

Each model has its own ZeroGPU pack (see zero_packs.py), and each button copies only the packs it
needs onto the GPU, one after another where a job uses several big models. That keeps every job
inside the default half-card slice, which is not billed double. If zero_packs is switched off or
cannot work, everything is copied on every request as ZeroGPU normally does, on the full card.
"""
import os


def _pick_model_store():
    """Off unless MODEL_STORE is set. Tried /data (the attached bucket) on 2026-10-04: it reads at
    about 8 MB/s, so loading 70 GB of models from it takes hours. The local disk downloads at 1 GB/s."""
    path = os.environ.get("MODEL_STORE")
    if not path:
        return None
    try:
        os.makedirs(os.path.join(path, "huggingface"), exist_ok=True)
        return os.path.join(path, "huggingface")
    except OSError:
        return None


_store = _pick_model_store()
if _store:
    os.environ["HF_HOME"] = _store          # Must be set before the Hugging Face libraries are imported

import json
import random
import re
import shutil
import subprocess
import sys
import threading
import time

import gradio as gr
import imageio
import imageio_ffmpeg
import numpy as np
import scipy.io.wavfile
import spaces
import torch
import zero_packs as packs
from diffusers import FluxPipeline, LTXConditionPipeline
from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
from diffusers.utils import export_to_video
from PIL import Image
from transformers import (AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor,
                          AutoTokenizer, MusicgenForConditionalGeneration, pipeline)
try:
    from transformers import AutoModelForVision2Seq
except ImportError:
    AutoModelForVision2Seq = None

IMAGE_MODEL = "black-forest-labs/FLUX.2-klein-9B"      # Gated, FLUX non-commercial licence
IMAGE_FALLBACK = "black-forest-labs/FLUX.1-schnell"    # Used if klein cannot be loaded (Apache 2.0)
VIDEO_MODEL = "Lightricks/LTX-Video-0.9.8-13B-distilled"
MUSIC_MODEL = "facebook/musicgen-medium"       # Non-commercial licence (CC-BY-NC 4.0)
TEXT_MODEL = "Qwen/Qwen3-4B"                   # Narration planner and prompt rewriter
IMAGE_CAPTION_MODEL = "florence-community/Florence-2-large"      # Picture to caption
VIDEO_CAPTION_MODEL = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"     # Video frames to description
LISTEN_MODEL = "laion/clap-htsat-unfused"      # Recognises genre, mood, instruments and tempo
SPEECH_MODEL = "openai/whisper-large-v3-turbo" # Speech to text
VOICE = "af_heart"                             # Kokoro voice
VOICE_RATE = 24000
# All 28 English voices of Kokoro. The first eight keep the names they had when there were only eight.
VOICES = {"Heart (woman)": "af_heart", "Bella (woman)": "af_bella", "Nicole (woman, soft)": "af_nicole",
          "Sarah (woman)": "af_sarah", "Adam (man)": "am_adam", "Michael (man)": "am_michael",
          "Fenrir (man, deep)": "am_fenrir", "Puck (man, lively)": "am_puck",
          "Alloy (woman)": "af_alloy", "Aoede (woman)": "af_aoede", "Jessica (woman)": "af_jessica",
          "Kore (woman)": "af_kore", "Nova (woman)": "af_nova", "River (woman)": "af_river", "Sky (woman)": "af_sky",
          "Echo (man)": "am_echo", "Eric (man)": "am_eric", "Liam (man)": "am_liam", "Onyx (man)": "am_onyx",
          "Santa (man)": "am_santa",
          "Alice (British woman)": "bf_alice", "Emma (British woman)": "bf_emma",
          "Isabella (British woman)": "bf_isabella", "Lily (British woman)": "bf_lily",
          "Daniel (British man)": "bm_daniel", "Fable (British man)": "bm_fable",
          "George (British man)": "bm_george", "Lewis (British man)": "bm_lewis"}
MAX_NARRATION_CHARS = 3000
WORDS_PER_SECOND = 2.5

IMAGE_WIDTH, IMAGE_HEIGHT = 1152, 768      # 3:2, multiples of 16 for FLUX
VIDEO_WIDTH, VIDEO_HEIGHT = 1152, 768      # Same size as the image, so nothing is shrunk (multiples of 32)
VIDEO_QUALITY = 9      # 0 to 10. Low values make the saved file blurry
FPS = 10
TRIM_FRAMES = 8        # Frames cut from the start of every clip, where the noise is
MAX_SEED = 2**31 - 1

NEGATIVE_PROMPT = ("worst quality, inconsistent motion, blurry, blurry motion, motion blur, "
                   "ignoring the prompt, jittery, distorted, watermark, text, deformed")
NEGATIVE_SLOW = ", slow motion, slow-motion video"
STYLE_SUFFIX = "Sharp, crisp, highly detailed, in focus, no blur"
SPEED_SUFFIX = ", moving at natural real-time speed"

# ZeroGPU attaches the GPU only inside @spaces.GPU calls. Each model is sent "to cuda" here and then
# sealed into its own pack, so a call copies only the packs it asks for with packs.load(...).
# Pure per-call .to('cuda') from ordinary RAM was tried on 2026-10-07: the first FLUX move took 79 s.
# The packs are read straight from disk at about 17 GB a second, like ZeroGPU's own single pack.
# Once a model is sealed, never call .to("cuda") on it again here: its weights are on disk.
print("[load] image model...", flush=True)
try:
    from diffusers import Flux2KleinPipeline
    image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_MODEL, torch_dtype=torch.bfloat16)
    IMAGE_IS_KLEIN = True
except Exception as e:      # Licence not accepted on the account, or diffusers too old: the Space still starts
    print(f"[load] {IMAGE_MODEL} failed ({type(e).__name__}: {str(e)[:200]}); using {IMAGE_FALLBACK}", flush=True)
    image_pipe = FluxPipeline.from_pretrained(IMAGE_FALLBACK, torch_dtype=torch.bfloat16)
    IMAGE_IS_KLEIN = False
for part in (image_pipe.transformer, image_pipe.vae, image_pipe.text_encoder,
             getattr(image_pipe, "text_encoder_2", None)):
    if part is not None:
        part.to("cuda")
packs.seal("image", IMAGE_MODEL if IMAGE_IS_KLEIN else IMAGE_FALLBACK)
print("[load] LTX...", flush=True)
video_pipe = LTXConditionPipeline.from_pretrained(VIDEO_MODEL, torch_dtype=torch.bfloat16)
video_pipe.vae.enable_tiling()     # Decode the video in tiles to keep memory low
video_pipe.text_encoder.to("cuda")
packs.seal("t5")
for part in (video_pipe.transformer, video_pipe.vae):
    part.to("cuda")
packs.seal("ltx", VIDEO_MODEL)
print("[load] MusicGen...", flush=True)
music_processor = AutoProcessor.from_pretrained(MUSIC_MODEL)
music_model = MusicgenForConditionalGeneration.from_pretrained(
    MUSIC_MODEL, torch_dtype=torch.float16).to("cuda")
packs.seal("musicgen", MUSIC_MODEL)
print("[load] Qwen3 4B...", flush=True)
text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL, torch_dtype=torch.bfloat16).to("cuda")
packs.seal("qwen", TEXT_MODEL)
print("[load] Florence-2...", flush=True)
image_caption_processor = AutoProcessor.from_pretrained(IMAGE_CAPTION_MODEL, trust_remote_code=True)
# florence-community is native; AutoModelForVision2Seq / AutoModelForImageTextToText register Florence2Config.
_florence_err = None
image_caption_model = None
for _cls in (AutoModelForVision2Seq, AutoModelForImageTextToText, AutoModelForCausalLM):
    if _cls is None:
        continue
    try:
        image_caption_model = _cls.from_pretrained(
            IMAGE_CAPTION_MODEL, torch_dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
        print(f"[load] Florence-2 via {_cls.__name__}", flush=True)
        break
    except Exception as e:
        _florence_err = e
        print(f"[load] Florence-2 via {_cls.__name__} failed: {type(e).__name__}: {e}", flush=True)
if image_caption_model is None:
    raise RuntimeError(f"Could not load Florence-2: {_florence_err}")
packs.seal("florence", IMAGE_CAPTION_MODEL)
print("[load] SmolVLM2...", flush=True)
video_caption_processor = AutoProcessor.from_pretrained(VIDEO_CAPTION_MODEL)
video_caption_model = AutoModelForImageTextToText.from_pretrained(
    VIDEO_CAPTION_MODEL, torch_dtype=torch.bfloat16).to("cuda")
packs.seal("smolvlm", VIDEO_CAPTION_MODEL)
print("[load] CLAP...", flush=True)
listen_pipe = pipeline("zero-shot-audio-classification", model=LISTEN_MODEL, device="cuda")
packs.seal("clap", LISTEN_MODEL)
print("[load] Whisper...", flush=True)
# No chunk_length_s: recordings longer than 30 seconds are then read with Whisper's own long-form
# method, one window after another. The pipeline's chunking was tried on 2026-10-11 and gave pieces
# with wrong times, repeated text and missing passages.
speech_pipe = pipeline("automatic-speech-recognition", model=SPEECH_MODEL, torch_dtype=torch.float16,
                       device="cuda")
packs.seal("whisper", SPEECH_MODEL)
print("[load] done.", flush=True)

# With a pack per model every job fits the default half-card slice ("large"), which is billed once.
# Without them ZeroGPU copies all the models on every call, and that needs the whole card.
# VIDEO_GPU_SIZE can be set to "xlarge" in the Space's variables if the video jobs run out of memory.
GPU_SIZE = "large" if packs.enabled else "xlarge"
VIDEO_GPU_SIZE = os.environ.get("VIDEO_GPU_SIZE", GPU_SIZE) if packs.enabled else "xlarge"


def gpu_seconds(seconds):
    """The GPU time a job asks for. The numbers passed in were measured on the half card on
    2026-10-08, where ZeroGPU counts 1.5 times what is asked (about 20% headroom is included).
    A job that runs past its time is cut off. Without per-model packs every job also waits about
    5 s for all the models to be copied, and the full card counts double."""
    return int(seconds) if packs.enabled else int(seconds + 6)


# ---------------------------------------------------------------- Timing helper

class timed:
    """Prints how long a block of code took."""
    def __init__(self, label):
        self.label = label

    def __enter__(self):
        self.t0 = time.time()
        return self

    def __exit__(self, *_):
        print(f"[{self.label}] {time.time() - self.t0:.2f}s", flush=True)


# ---------------------------------------------------------------- Kokoro (runs on the CPU in its own process)

def start_voice_worker():
    worker = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tts_worker.py")
    return subprocess.Popen([sys.executable, worker], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)


voice_lock = threading.Lock()
voice_worker = start_voice_worker()


def kokoro(text, voice, speed=1.0):
    """The text spoken by Kokoro, as a float32 wave at VOICE_RATE."""
    global voice_worker
    with voice_lock:
        if voice_worker.poll() is not None:
            voice_worker = start_voice_worker()
        out = f"/tmp/kokoro_{random.randint(0, MAX_SEED)}.npy"
        voice_worker.stdin.write(json.dumps({"text": text, "voice": voice, "speed": speed, "out": out}) + "\n")
        voice_worker.stdin.flush()
        answer = voice_worker.stdout.readline()
        if not answer:
            raise gr.Error("The narration model stopped. Please try again.")
        answer = json.loads(answer)
        if not answer["ok"]:
            raise gr.Error("Narration failed: " + answer["error"])
        wave = np.load(out)
        os.remove(out)
        return wave


def pick_seed(seed, randomize):
    return random.randint(0, MAX_SEED) if randomize else int(seed)


# ---------------------------------------------------------------- The core model steps (run on the GPU)

def draw_image(prompt, seed):
    # Both are 4-step distilled models; they only differ in how guidance is switched off
    options = {"guidance_scale": 1.0} if IMAGE_IS_KLEIN else {"guidance_scale": 0.0, "max_sequence_length": 256}
    return image_pipe(
        prompt=prompt,
        width=IMAGE_WIDTH,
        height=IMAGE_HEIGHT,
        num_inference_steps=4,
        generator=torch.Generator("cpu").manual_seed(seed),
        **options,
    ).images[0]


def animate(image, scene, motion, seconds, seed):
    """Returns the path of a silent clip and its length in seconds."""
    num_frames = clip_frames(seconds)
    motion = motion.strip() or "gentle cinematic camera motion, natural movement"
    wants_slow = re.search(r"slow[\s-]?mo", motion.lower()) is not None
    style = STYLE_SUFFIX + ("" if wants_slow else SPEED_SUFFIX) + "."
    prompt = f"{motion}. {scene.strip()}. {style}" if scene.strip() else f"{motion}. {style}"
    start = LTXVideoCondition(image=image.convert("RGB").resize((VIDEO_WIDTH, VIDEO_HEIGHT)), frame_index=0)
    frames = video_pipe(
        conditions=[start],
        prompt=prompt,
        negative_prompt=NEGATIVE_PROMPT + ("" if wants_slow else NEGATIVE_SLOW),
        image_cond_noise_scale=0.0,
        decode_timestep=0.05,
        decode_noise_scale=0.025,
        width=VIDEO_WIDTH,
        height=VIDEO_HEIGHT,
        num_frames=num_frames,
        frame_rate=FPS,
        timesteps=[1000, 993, 987, 981, 975, 909, 725, 0.03],
        guidance_scale=1.0,
        guidance_rescale=0.7,
        generator=torch.Generator("cpu").manual_seed(seed),
    ).frames[0]
    frames = frames[TRIM_FRAMES:]
    path = f"/tmp/video_{seed}_{random.randint(0, 99999)}.mp4"
    export_to_video(frames, path, fps=FPS, quality=VIDEO_QUALITY)
    return path, len(frames) / FPS


def compose_music(music_prompt, seconds):
    """Returns the path of a WAV file as long as the clip."""
    seconds = min(max(seconds, 2.0), 28.0)
    inputs = music_processor(text=[music_prompt], padding=True, return_tensors="pt").to("cuda")
    audio = music_model.generate(**inputs, do_sample=True, guidance_scale=3.0,
                                 max_new_tokens=int(seconds * 50) + 8)
    rate = music_model.config.audio_encoder.sampling_rate
    wave = audio[0, 0].float().cpu().numpy()[:int(seconds * rate)]
    fade = int(0.25 * rate)
    wave[:fade] *= np.linspace(0.0, 1.0, fade)
    wave[-fade:] *= np.linspace(1.0, 0.0, fade)
    wav_path = f"/tmp/music_{random.randint(0, MAX_SEED)}.wav"
    scipy.io.wavfile.write(wav_path, rate, (np.clip(wave, -1, 1) * 32767).astype(np.int16))
    return wav_path


# ---------------------------------------------------------------- Describing (vision / audio) and rewriting (text)

GENRES = ["classical orchestral", "electronic dance", "rock", "jazz", "hip hop", "ambient", "pop", "folk acoustic",
          "8-bit chiptune", "cinematic soundtrack", "lo-fi", "heavy metal", "Indian classical", "devotional chant"]
MOODS = ["happy and upbeat", "calm and peaceful", "sad and emotional", "tense and dramatic", "epic and heroic",
         "dark and mysterious", "playful and fun", "dreamy and floating"]
INSTRUMENTS = ["piano", "acoustic guitar", "electric guitar", "violins and strings", "drums", "synthesizer", "flute",
               "bass", "brass and trumpets", "singing voice", "bells", "sitar", "tabla"]
TEMPOS = ["slow tempo", "medium tempo", "fast tempo"]


# Chat models like to open with "Certainly! Here's a revised prompt:". That opening is cut off,
# because the image, video and music models would otherwise be given it as part of the prompt.
OPENING_WORD = re.compile(r"^(certainly|sure|of course|absolutely|okay|ok|great)\b[!,.]?\s*", re.IGNORECASE)
OPENING_LINE = re.compile(r"^(here(['’]s| is| are)|below is|this is)\b[^\n.!?]*:\s*", re.IGNORECASE)


def tidy(text, one_line=False):
    """The model's answer without its chatty first sentence and without Markdown marks."""
    text = text.strip()
    text = OPENING_WORD.sub("", text, count=1)
    text = OPENING_LINE.sub("", text, count=1)
    text = re.sub(r"^#+\s*", "", text, flags=re.MULTILINE)
    text = re.sub(r"[*`]+", "", text).strip().strip('"').strip()
    return " ".join(text.split()) if one_line else text


def shrink(image, side=768):
    image = image.convert("RGB")
    image.thumbnail((side, side))
    return image


def caption_image(image):
    """Florence-2 writes a detailed caption. Model must already be on the GPU."""
    prompt = "<MORE_DETAILED_CAPTION>"
    inputs = image_caption_processor(text=prompt, images=image, return_tensors="pt").to("cuda", torch.bfloat16)
    out = image_caption_model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"],
        max_new_tokens=512,
        num_beams=3,
        do_sample=False,
    )
    decoded = image_caption_processor.batch_decode(out, skip_special_tokens=False)[0]
    parsed = image_caption_processor.post_process_generation(
        decoded, task=prompt, image_size=(image.width, image.height))
    return parsed.get(prompt, "").strip()


def caption_video_frames(frames):
    """SmolVLM2 describes a short video from a few frames. Model must already be on the GPU."""
    content = [{"type": "image", "image": f} for f in frames]
    content.append({"type": "text", "text": "These pictures are frames from one video, in order. "
                                             "Describe the scene, then describe what moves and how."})
    inputs = video_caption_processor.apply_chat_template(
        [{"role": "user", "content": content}],
        add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt").to("cuda")
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)
    out = video_caption_model.generate(**inputs, max_new_tokens=300, do_sample=False)
    text = video_caption_processor.batch_decode(
        out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
    return tidy(text)


def caption_music(wave):
    """CLAP picks a genre, mood, tempo and two instruments, written as one short sentence.
    Model must already be on the GPU."""
    def best(labels, count=1):
        return [hit["label"] for hit in listen_pipe(wave, candidate_labels=labels)[:count]]
    return (f"{best(GENRES)[0]} music, {best(MOODS)[0]}, {best(TEMPOS)[0]}, "
            f"featuring {' and '.join(best(INSTRUMENTS, 2))}")


def rewrite(system, user, max_new_tokens=160, one_line=True):
    """Qwen3 4B edits a description into a better prompt. Model must already be on the GPU."""
    text = text_tokenizer.apply_chat_template(
        [{"role": "system", "content": system}, {"role": "user", "content": user}],
        tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = text_tokenizer([text], return_tensors="pt").to("cuda")
    out = text_model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True,
                              temperature=0.6, top_p=0.9)
    return tidy(text_tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True), one_line)


def plan_narration(scene, motion, seconds):
    """Qwen3 4B writes 1 to 3 lines of narration with start times. Model must already be on the GPU."""
    budget = max(4, int((seconds - 1) * WORDS_PER_SECOND))
    system = ("You are the narrator of a short film clip. You are given what the clip shows. Write what the "
              "narrator SAYS over it: a thought, a feeling or a tiny story about the scene, like a nature "
              "documentary. Never repeat or list the description itself, and never mention cameras, prompts or "
              "video. Reply with 1 to 3 lines and nothing else. Each line is: start second | sentence.\n"
              f"The clip lasts {seconds:.0f} seconds. Use at most {budget} words in total, so every line is short. "
              "The first line starts at 0.5 or later, and later lines start after the earlier ones have been spoken.\n\n"
              "Example for a clip of a lighthouse in a storm, 6 seconds:\n"
              "0.5 | Some lights refuse to go out.\n"
              "3.5 | Even when the sea says otherwise.")
    user = f"The clip shows: {scene or 'a quiet scene'}\nWhat moves: {motion or 'gentle movement'}"
    text = text_tokenizer.apply_chat_template(
        [{"role": "system", "content": system}, {"role": "user", "content": user}],
        tokenize=False, add_generation_prompt=True, enable_thinking=False)
    inputs = text_tokenizer([text], return_tensors="pt").to("cuda")
    out = text_model.generate(**inputs, max_new_tokens=120, do_sample=True, temperature=0.8, top_p=0.9)
    return text_tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)


# ---------------------------------------------------------------- Narration and mixing (CPU only)

def parse_plan(reply, seconds):
    """Turn the planner's reply into [(start, text)]. If it ignored the format, its own words are
    still used as one line; the scene description is never read out."""
    plan = []
    for row in reply.splitlines():
        match = re.match(r"\s*[-*]?\s*(\d+(?:\.\d+)?)\s*(?:s|sec|seconds)?\s*[|:–-]\s*(.+)", row)
        if match and match.group(2).strip():
            plan.append((float(match.group(1)), match.group(2).strip().strip('"')))
    if not plan:
        words = re.sub(r"[\[\]{}\"|]", " ", reply).split()[:int(seconds * WORDS_PER_SECOND)]
        plan = [(0.5, " ".join(words) or "A quiet moment, captured in motion.")]
    return sorted(plan)[:3]


def speak(text):
    return kokoro(text, VOICE)


def narrate(reply, seconds):
    """Speak the plan into one WAV as long as the clip. Returns (wav path or None, the plan as text)."""
    track = np.zeros(int(seconds * VOICE_RATE), dtype=np.float32)
    cursor, script = 0.0, []
    for start, line in parse_plan(reply, seconds):
        voice = speak(line)
        start = max(start, cursor)
        begin = int(start * VOICE_RATE)
        room = len(track) - begin
        if room < VOICE_RATE // 2:
            continue
        voice = voice[:room]
        track[begin:begin + len(voice)] += voice
        cursor = start + len(voice) / VOICE_RATE + 0.3
        script.append(f"{start:4.1f} s  {line}")
    if not script:
        return None, "The clip is too short for narration."
    path = f"/tmp/voice_{random.randint(0, MAX_SEED)}.wav"
    scipy.io.wavfile.write(path, VOICE_RATE, (np.clip(track, -1, 1) * 32767).astype(np.int16))
    return path, "\n".join(script)


def mix(video, music_wav, voice_wav):
    """Put whichever sound tracks exist onto the silent clip. With both, music plays under the voice."""
    if not music_wav and not voice_wav:
        return video
    ffmpeg = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-i", video]
    if music_wav and voice_wav:
        ffmpeg += ["-i", voice_wav, "-i", music_wav, "-filter_complex",
                   "[1:a]aresample=44100,volume=1.6[v];[2:a]aresample=44100,volume=0.35[m];"
                   "[v][m]amix=inputs=2:duration=longest:normalize=0[a]", "-map", "0:v:0", "-map", "[a]"]
    else:
        ffmpeg += ["-i", music_wav or voice_wav, "-map", "0:v:0", "-map", "1:a:0"]
    out_path = f"/tmp/mixed_{random.randint(0, MAX_SEED)}.mp4"
    subprocess.run(ffmpeg + ["-c:v", "copy", "-c:a", "aac", "-shortest", out_path], check=True, capture_output=True)
    return out_path


def load_audio(path, limit=30):
    """Any audio file as mono 48 kHz samples (the first 30 seconds), using ffmpeg."""
    wav = f"/tmp/listen_{random.randint(0, MAX_SEED)}.wav"
    subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-i", path, "-t", str(limit), "-ac", "1", "-ar", "48000", wav],
                   check=True, capture_output=True)
    rate, data = scipy.io.wavfile.read(wav)
    return data.astype(np.float32) / 32768.0


def video_frames(path, count=6):
    """A few frames spread across the clip, and its length in seconds."""
    reader = imageio.get_reader(path)
    meta = reader.get_meta_data()
    length = float(meta.get("duration") or 5.0)
    total = max(1, int(length * float(meta.get("fps") or 10)) - 1)
    frames = []
    for k in range(count):
        try:
            frames.append(shrink(Image.fromarray(reader.get_data(total * k // max(1, count - 1))), 512))
        except Exception:
            break
    reader.close()
    return frames, length


# ---------------------------------------------------------------- The buttons

@spaces.GPU(duration=gpu_seconds(7), size=GPU_SIZE)
def make_image(prompt, seed, randomize_seed):
    if not prompt.strip():
        raise gr.Error("Describe the image first.")
    seed = pick_seed(seed, randomize_seed)
    with timed("image"):
        packs.load("image")
        return draw_image(prompt, seed), seed


def music_duration(prompt, seconds):
    return gpu_seconds(3 + seconds * 1.0)


@spaces.GPU(duration=music_duration, size=GPU_SIZE)
def make_music(prompt, seconds):
    if not prompt.strip():
        raise gr.Error("Describe the music first.")
    with timed("music"):
        packs.load("musicgen")
        return compose_music(prompt, seconds)


def audio_seconds(path):
    """Length of an audio file in seconds, read from ffmpeg's report. 30 if it cannot be read."""
    try:
        report = subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-hide_banner", "-i", path],
                                capture_output=True, text=True, timeout=30).stderr
        hours, minutes, seconds = re.search(r"Duration: (\d+):(\d+):(\d+\.?\d*)", report).groups()
        return int(hours) * 3600 + int(minutes) * 60 + float(seconds)
    except Exception:
        return 30.0


def transcribe_duration(audio, task):
    # A short clip needs about 3 s in all. Long recordings are read in 30-second pieces; this allows
    # for the model working at 8 times real speed, which is slower than it has measured so far.
    return gpu_seconds(5 + audio_seconds(audio) / 12) if audio else gpu_seconds(5)


@spaces.GPU(duration=transcribe_duration, size=GPU_SIZE)
def transcribe(audio, task):
    if not audio:
        raise gr.Error("Upload a recording or record with the microphone first.")
    with timed("transcribe"):
        packs.load("whisper")
        options = {"task": "translate" if task.startswith("Translate") else "transcribe"}
        if audio_seconds(audio) > 30:       # Long-form only: retry a window that came out badly, skip silence
            options.update(temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), no_speech_threshold=0.6,
                           logprob_threshold=-1.0, compression_ratio_threshold=1.35)
        result = speech_pipe(audio, return_timestamps=True, generate_kwargs=options)
    # The same text again, one piece per line with its start and end time in seconds
    timed_lines = []
    for piece in result.get("chunks") or []:
        start, end = piece.get("timestamp") or (None, None)
        if start is not None and piece.get("text", "").strip():
            timed_lines.append(f"[{start:.2f} - {end if end is not None else start:.2f}] {piece['text'].strip()}")
    return result["text"].strip(), "\n".join(timed_lines)


def read_aloud(text, voice_name, speed):
    """Kokoro runs on the CPU, so this uses no GPU quota."""
    text = text.strip()
    if not text:
        raise gr.Error("Type the text to read first.")
    if len(text) > MAX_NARRATION_CHARS:
        raise gr.Error(f"That is {len(text)} characters. Please keep it under {MAX_NARRATION_CHARS}.")
    wave = kokoro(text, VOICES[voice_name], float(speed))
    if not len(wave):
        raise gr.Error("Nothing could be read from that text.")
    path = f"/tmp/narration_{random.randint(0, MAX_SEED)}.wav"
    scipy.io.wavfile.write(path, VOICE_RATE, (np.clip(wave, -1, 1) * 32767).astype(np.int16))
    return path


# ---------------------------------------------------------------- Describe tab

@spaces.GPU(duration=gpu_seconds(5), size=GPU_SIZE)
def describe_image_tab(image):
    if image is None:
        raise gr.Error("Upload an image first.")
    small = shrink(image)
    with timed("image-to-text"):
        packs.load("florence")
        return caption_image(small)


@spaces.GPU(duration=gpu_seconds(6), size=GPU_SIZE)
def describe_video_tab(video):
    if not video:
        raise gr.Error("Upload a video first.")
    frames, _ = video_frames(video)
    if not frames:
        raise gr.Error("That video could not be read.")
    with timed("video-to-text"):
        packs.load("smolvlm")
        return caption_video_frames(frames)


@spaces.GPU(duration=gpu_seconds(4), size=GPU_SIZE)
def describe_music_tab(audio):
    if not audio:
        raise gr.Error("Upload some music first.")
    try:
        wave = load_audio(audio)
    except Exception:
        raise gr.Error("That audio file could not be read.")
    with timed("music-to-text"):
        packs.load("clap")
        return caption_music(wave)


# ---------------------------------------------------------------- Video tab

def improve_prompts(scene, motion):
    """Qwen3 4B adds visual detail to the scene and motion prompts without changing what they ask for."""
    reply = rewrite(
        "You improve prompts for a video generator. You are given a scene and a motion. Keep every subject, "
        "action and direction exactly as given (left stays left, right stays right), and never add new subjects "
        "or change what moves. Add what a camera would see: lighting, colours, textures, atmosphere, depth. "
        "Reply with exactly two lines and nothing else:\nSCENE: the picture, under 60 words\n"
        "MOTION: what moves and how, under 30 words",
        f"Scene: {scene.strip() or 'not given, keep it short and neutral'}\n"
        f"Motion: {motion.strip() or 'gentle natural movement'}", 200, one_line=False)
    # The model answers in several shapes: "SCENE: ... MOTION: ..." on one line or two, or two bare lines
    label = re.search(r"\bmotion\s*:\s*", reply, re.IGNORECASE)
    lines = [line.strip() for line in reply.splitlines() if line.strip()]
    if label:
        new_scene, new_motion = reply[:label.start()], reply[label.end():]
    elif len(lines) == 2:
        new_scene, new_motion = lines
    else:
        new_scene = new_motion = ""
    new_scene = " ".join(re.sub(r"^\s*scene\s*:\s*", "", new_scene, flags=re.IGNORECASE).split())
    new_motion = " ".join(new_motion.split())
    if not new_scene or not new_motion:                 # Anything else: keep the visitor's own words
        print(f"[improve] unusable answer: {reply[:200]!r}", flush=True)
        return scene, motion
    return (new_scene if scene.strip() else scene), new_motion


def clip_frames(seconds):
    """Frames LTX-Video is asked for: a multiple of 8 plus 1, with room for the trimmed start."""
    return -(-(int(seconds * FPS) + TRIM_FRAMES - 1) // 8) * 8 + 1


def film(image, scene, motion, music_prompt, want_narration, improve, seconds, seed):
    """A clip made one model at a time, so FLUX and LTX-Video are never on the GPU together:
    the text model (prompts and narration plan), FLUX draws, LTX-Video animates, then MusicGen."""
    plan = None
    improve = improve and bool(scene.strip() or motion.strip())
    if improve or want_narration:
        packs.load("qwen")
        with timed("  text"):
            if improve:
                scene, motion = improve_prompts(scene, motion)
            if want_narration:
                plan = plan_narration(scene, motion, (clip_frames(seconds) - TRIM_FRAMES) / FPS)
    if image is None:
        packs.load("image")
        with timed("  draw"):
            image = draw_image(scene, seed)
    packs.load("t5", "ltx")
    with timed("  animate"):
        video, length = animate(image, scene, motion, seconds, seed)
    music_wav = None
    if music_prompt.strip():
        packs.load("musicgen")
        with timed("  music"):
            music_wav = compose_music(music_prompt, length)
    return video, length, music_wav, plan, scene, motion


def studio_duration(image, scene, motion, music_prompt, want_narration, improve, seconds, seed):
    total = 7 + seconds * 3.0                       # Animation
    if image is None:
        total += 5                                  # FLUX draws the scene first
    if music_prompt.strip():
        total += 2 + seconds * 1.0
    if want_narration:
        total += 1
    if improve:
        total += 2                                  # The text model rewrites the prompts
    return gpu_seconds(total)


@spaces.GPU(duration=studio_duration, size=VIDEO_GPU_SIZE)
def run_models(image, scene, motion, music_prompt, want_narration, improve, seconds, seed):
    with timed("video"):
        return film(image, scene, motion, music_prompt, want_narration, improve, seconds, seed)


def make_video(image, scene, motion, music_prompt, want_narration, improve, seconds, seed, randomize_seed):
    if image is None and not scene.strip():
        raise gr.Error("Describe the scene, or add a start image.")
    seed = pick_seed(seed, randomize_seed)
    video, length, music_wav, plan, used_scene, used_motion = run_models(
        image, scene, motion, music_prompt, want_narration, improve, seconds, seed)
    voice_wav, script = narrate(plan, length) if plan is not None else (None, "")
    used = f"SCENE: {used_scene}\nMOTION: {used_motion}" if improve else ""
    return mix(video, music_wav, voice_wav), script, used, seed


# ---------------------------------------------------------------- Bucket: running totals and showcase

BUCKET = os.environ.get("STUDIO_BUCKET", "/data")
COUNTS_FILE = os.path.join(BUCKET, "counts.json")
SHOWCASE_DIR = os.path.join(BUCKET, "showcase")
SHOWCASE_COPY = "/tmp/showcase"
SHOWCASE_TYPES = (".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm")
COUNT_NAMES = {"images": "images", "videos": "videos", "songs": "songs", "narrations": "narrations"}
counts_lock = threading.Lock()
counts = dict.fromkeys(COUNT_NAMES, 0)
try:
    with open(COUNTS_FILE, encoding="utf-8") as f:
        counts.update({k: int(v) for k, v in json.load(f).items() if k in counts})
except (OSError, ValueError, AttributeError, TypeError):
    pass


def counts_text():
    total = sum(counts.values())
    if not total:
        return ""
    parts = " · ".join(f"{counts[k]:,} {COUNT_NAMES[k][:-1] if counts[k] == 1 else COUNT_NAMES[k]}"
                       for k in COUNT_NAMES if counts[k])
    return f"**{total:,} creation{'' if total == 1 else 's'} so far:** {parts}"


def bump(kind):
    with counts_lock:
        counts[kind] += 1
        try:
            with open(COUNTS_FILE, "w", encoding="utf-8") as f:
                json.dump(counts, f)
        except OSError:
            pass
    return counts_text()


def showcase():
    try:
        names = sorted(n for n in os.listdir(SHOWCASE_DIR) if n.lower().endswith(SHOWCASE_TYPES))
    except OSError:
        return []
    os.makedirs(SHOWCASE_COPY, exist_ok=True)
    shown = []
    for name in names:
        local = os.path.join(SHOWCASE_COPY, name)
        try:
            source = os.path.join(SHOWCASE_DIR, name)
            if not os.path.exists(local) or os.path.getsize(local) != os.path.getsize(source):     # New, or replaced
                shutil.copy(source, local + ".part")
                os.replace(local + ".part", local)
            shown.append((local, os.path.splitext(name)[0].replace("_", " ")))
        except OSError:
            continue
    return shown


def on_open():
    return counts_text(), showcase()


# ---------------------------------------------------------------- UI

with gr.Blocks(title="Graphics Studio") as demo:
    gr.Markdown("# Graphics Studio\n**FLUX.2 klein** images, **LTX-Video** animation, **MusicGen** music, "
                "**Qwen3** + **Kokoro** narration, **Whisper** transcription, "
                "**Florence-2** image captions, **SmolVLM2** video descriptions, **CLAP** music classification.")
    gr.Markdown("Already have an image, a video or some music and want to change it? Try my other Space, "
                "[Graphics Enhancer](https://huggingface.co/spaces/vrishv/graphics_enhancer).")
    counter = gr.Markdown(counts_text())

    with gr.Tab("Image"):
        image_prompt = gr.Textbox(label="Describe the image", lines=3,
                                  placeholder="A breathtaking aurora over a frozen lake at night, ultra detailed")
        with gr.Row():
            image_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
            image_random = gr.Checkbox(value=True, label="Random seed")
        image_button = gr.Button("Generate image", variant="primary")
        image_out = gr.Image(label="Image", type="pil")
        gr.Examples(
            examples=[
                ["A breathtaking aurora over a frozen lake at night, snowy mountains, ultra detailed"],
                ["A lone lighthouse on a cliff during a thunderstorm, giant waves, cinematic lighting"],
                ["A floating island with waterfalls at sunrise, golden clouds, fantasy landscape"],
            ],
            inputs=[image_prompt],
        )

    with gr.Tab("Music"):
        song_prompt = gr.Textbox(label="Describe the music", lines=3,
                                 placeholder="Upbeat 8-bit arcade tune with a catchy melody, fast tempo")
        song_seconds = gr.Slider(5, 28, value=15, step=1, label="Length (seconds)")
        song_button = gr.Button("Generate music", variant="primary")
        song_out = gr.Audio(label="Music", type="filepath", autoplay=False)
        gr.Examples(
            examples=[
                ["Upbeat 8-bit arcade tune with a catchy melody, fast tempo"],
                ["Calm ambient synth pads, slow and dreamy, night sky mood"],
                ["Epic orchestral trailer music with drums and brass, building tension"],
            ],
            inputs=[song_prompt],
        )

    with gr.Tab("Transcribe"):
        speech_in = gr.Audio(label="Upload a recording, or speak now", sources=["upload", "microphone"],
                             type="filepath")
        speech_task = gr.Radio(["Transcribe (keep the spoken language)", "Translate to English"],
                               value="Transcribe (keep the spoken language)", label="What to write")
        speech_button = gr.Button("Transcribe", variant="primary")
        speech_out = gr.Textbox(label="Text", lines=8)
        speech_times = gr.Textbox(label="With times (start and end of each piece, in seconds)", lines=8,
                                  interactive=False)

    with gr.Tab("Narration"):
        read_text = gr.Textbox(label="Text to read aloud", lines=6,
                               placeholder="Far above the sleeping town, the stars kept their quiet watch.")
        with gr.Row():
            read_voice = gr.Dropdown(list(VOICES), value="Heart (woman)", label="Voice")
            read_speed = gr.Slider(0.7, 1.3, value=1.0, step=0.05, label="Speed")
        read_button = gr.Button("Read aloud", variant="primary")
        read_out = gr.Audio(label="Narration", type="filepath", autoplay=False)

    with gr.Tab("Describe"):
        gr.Markdown("Three ways to describe what you have: an image (image captioning), a short video "
                    "(video captioning), or a piece of music (music captioning, music classification). "
                    "The video and music jobs look at the first 30 seconds.")
        with gr.Tabs():
            with gr.Tab("Image to text"):
                with gr.Row():
                    with gr.Column():
                        desc_image_in = gr.Image(label="Your image", type="pil")
                        desc_image_button = gr.Button("Describe image", variant="primary")
                    with gr.Column():
                        desc_image_out = gr.Textbox(label="Description", lines=10, interactive=False)
            with gr.Tab("Video to text"):
                with gr.Row():
                    with gr.Column():
                        desc_video_in = gr.Video(label="Your video", autoplay=False)
                        desc_video_button = gr.Button("Describe video", variant="primary")
                    with gr.Column():
                        desc_video_out = gr.Textbox(label="Description", lines=10, interactive=False)
            with gr.Tab("Music to text"):
                with gr.Row():
                    with gr.Column():
                        desc_music_in = gr.Audio(label="Your music", sources=["upload"], type="filepath")
                        desc_music_button = gr.Button("Describe music", variant="primary")
                    with gr.Column():
                        desc_music_out = gr.Textbox(label="Description", lines=4, interactive=False)

    with gr.Tab("Video"):
        with gr.Row():
            with gr.Column():
                scene_prompt = gr.Textbox(label="Scene", lines=3,
                                          placeholder="A lone lighthouse on a cliff during a thunderstorm, giant waves")
                motion_prompt = gr.Textbox(label="Motion", lines=2,
                                           placeholder="Waves crash against the rocks, the light beam sweeps to the right")
                music_prompt = gr.Textbox(label="Music (optional, leave empty for none)", lines=2,
                                          placeholder="Tense orchestral strings with distant thunder")
                improve_prompt = gr.Checkbox(value=True, label="Improve my prompts first (Qwen3 4B adds visual "
                                                                "detail; untick to use your exact words)")
                want_narration = gr.Checkbox(value=False, label="Add narration")
                seconds = gr.Slider(2, 8, value=5, step=1, label="Length (seconds)")
                with gr.Accordion("More options", open=False):
                    start_image = gr.Image(label="Start image (optional; otherwise FLUX draws the scene)", type="pil")
                    with gr.Row():
                        video_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        video_random = gr.Checkbox(value=True, label="Random seed")
                video_button = gr.Button("Make video", variant="primary")
            with gr.Column():
                video_out = gr.Video(label="Video", autoplay=False, loop=True)
                narration_script = gr.Textbox(label="Narration (start time and line)", lines=3, interactive=False)
                used_prompts = gr.Textbox(label="Prompts used (after improving)", lines=4, interactive=False)

    with gr.Tab("Showcase"):
        gr.Markdown("A few favourites made with this Space.")
        showcase_out = gr.Gallery(label="Showcase", columns=3)

    image_button.click(make_image, [image_prompt, image_seed, image_random], [image_out, image_seed]).success(
        lambda: bump("images"), None, counter)
    song_button.click(make_music, [song_prompt, song_seconds], [song_out]).success(
        lambda: bump("songs"), None, counter)
    speech_button.click(transcribe, [speech_in, speech_task], [speech_out, speech_times])
    read_button.click(read_aloud, [read_text, read_voice, read_speed], [read_out]).success(
        lambda: bump("narrations"), None, counter)
    desc_image_button.click(describe_image_tab, [desc_image_in], [desc_image_out])
    desc_video_button.click(describe_video_tab, [desc_video_in], [desc_video_out])
    desc_music_button.click(describe_music_tab, [desc_music_in], [desc_music_out])
    video_button.click(make_video,
                       [start_image, scene_prompt, motion_prompt, music_prompt, want_narration, improve_prompt,
                        seconds, video_seed, video_random],
                       [video_out, narration_script, used_prompts, video_seed]).success(
        lambda: bump("videos"), None, counter)
    demo.load(on_open, None, [counter, showcase_out])

demo.queue().launch(allowed_paths=[SHOWCASE_COPY])