Spaces:

Beexly
/

studio-media

Paused

Upload app.py with huggingface_hub
2641ca9 verified 2 months ago
32.3 kB
import json
import os
import random
import tempfile
import uuid

import numpy as np
import soundfile as sf
import spaces
import torch
import gradio as gr
from diffusers import FluxPipeline
from gradio_client import Client as GradioClient
from transformers import AutoModelForImageSegmentation, AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from kokoro import KPipeline
from f5_tts.api import F5TTS

# ---------------------------------------------------------------------------
# Fix a real bug in gradio_client 1.5.2 (bundled with ZeroGPU's forced
# gradio==5.9.1): its JSON-schema-to-python-type parser assumes every schema
# node is a dict, but valid JSON Schema allows a bare boolean (e.g.
# additionalProperties: true) — that crashes EVERY API call to this Space
# with "TypeError: argument of type 'bool' is not iterable". Patched here,
# not worked around with an old pydantic pin (which broke docling's own
# serialization elsewhere) — this is the actual bug.
import gradio_client.utils as _gc_utils

_orig_json_schema_to_python_type = _gc_utils._json_schema_to_python_type


def _patched_json_schema_to_python_type(schema, defs):
    if not isinstance(schema, dict):
        return "bool" if isinstance(schema, bool) else "Any"
    return _orig_json_schema_to_python_type(schema, defs)


_gc_utils._json_schema_to_python_type = _patched_json_schema_to_python_type

# ---------------------------------------------------------------------------
# basicsr 1.4.2 (a realesrgan dependency, used by the Upscale tab) imports
# torchvision.transforms.functional_tensor, removed in torchvision>=0.17 (this
# Space has 0.28.0) — real ModuleNotFoundError from a live call. Shim it back
# with the one function basicsr actually needs, before basicsr is imported.
import sys
import types
import torchvision.transforms.functional as _tvf

_functional_tensor_shim = types.ModuleType("torchvision.transforms.functional_tensor")
_functional_tensor_shim.rgb_to_grayscale = _tvf.rgb_to_grayscale
sys.modules["torchvision.transforms.functional_tensor"] = _functional_tensor_shim

# ---------------------------------------------------------------------------
# Models — loaded once at import time. Under ZeroGPU's patched torch, .to("cuda")
# / device_map="auto" here is deferred (no real GPU touched yet); the ZeroGPU
# "tensors packing" step (visible in build logs) prepares each model for fast
# real materialization inside its own @spaces.GPU-decorated call.
# ---------------------------------------------------------------------------

IMAGE_MODELS = {
    "Best quality (FLUX.1-dev)": "black-forest-labs/FLUX.1-dev",
}
RESOLUTIONS = {
    "1024x1024 (square)": (1024, 1024),
    "1152x896 (landscape)": (1152, 896),
    "896x1152 (portrait)": (896, 1152),
    "1344x768 (wide)": (1344, 768),
    "768x1344 (tall)": (768, 1344),
}
ASR_MODEL_ID = "openai/whisper-large-v3-turbo"
TTS_VOICES = ["af_heart", "af_bella", "am_michael", "bf_emma"]
STUDIO_CHAT_SPACE = "Beexly/studio-chat"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Primary: FLUX.1-dev, the proven deferred-cuda pattern (packed at Space startup).
image_pipe = FluxPipeline.from_pretrained(IMAGE_MODELS["Best quality (FLUX.1-dev)"], torch_dtype=torch.bfloat16)
image_pipe.to("cuda")

# FLUX.1-schnell ("fast draft") was tried as a second, lazily-GPU-loaded
# pipeline and reverted — see README "Left out: FLUX.1-schnell fast draft"
# for the real VRAM-capacity math (confirmed via Space logs, not a guess).

# AntiBlur LoRA (Shakker-Labs, same FLUX.1-dev non-commercial license as the
# base model — no new restriction) — deepens depth-of-field / reduces
# background blur bleed. Loaded once as a non-destructive adapter (NOT the
# model card's pipe.fuse_lora(), which permanently bakes it into every
# future generation on this shared pipeline) — toggled on/off per request
# via set_adapters()/disable_lora() so it stays fully reversible.
image_pipe.load_lora_weights(
    "Shakker-Labs/FLUX.1-dev-LoRA-AntiBlur",
    weight_name="FLUX-dev-lora-AntiBlur.safetensors",
    adapter_name="antiblur",
)
# disable_lora() + set_adapters() to re-enable is a documented diffusers/PEFT
# footgun (github.com/huggingface/diffusers/issues/11148 and others) — a
# real call with this Space's first version produced byte-identical output
# with the toggle on vs off, confirming it silently no-ops here too. Fixed
# by never disabling the adapter at all: always call set_adapters() with an
# explicit weight, 0.0 to turn it off, avoiding that code path entirely.
image_pipe.set_adapters(["antiblur"], adapter_weights=[0.0])

asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(
    ASR_MODEL_ID, torch_dtype=torch.float16, low_cpu_mem_usage=True, use_safetensors=True
)
asr_model.to("cuda")
asr_processor = AutoProcessor.from_pretrained(ASR_MODEL_ID)
asr_pipe = pipeline(
    "automatic-speech-recognition",
    model=asr_model,
    tokenizer=asr_processor.tokenizer,
    feature_extractor=asr_processor.feature_extractor,
    torch_dtype=torch.float16,
    device="cuda",
)

kokoro_pipeline = KPipeline(lang_code="a")

# BiRefNet (MIT, ZhengPeng7/BiRefNet) — 0.2B params, small enough to pack
# alongside the models above rather than lazy-load.
birefnet_model = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet_model.to("cuda")
birefnet_model.eval()

# F5TTS's own loader calls safetensors.load_file(..., device="cuda") directly,
# which needs a real GPU *at import time* — incompatible with ZeroGPU's normal
# deferred-.to("cuda") pattern. Load on CPU here; move to CUDA lazily inside
# the @spaces.GPU-decorated function that actually uses it.
f5tts_model = F5TTS(device="cpu")
_f5tts_on_gpu = False

# MusicGen-small (facebook/musicgen-small) — transformers-native, clean
# documented API, small (0.6B) — lazy CPU-load-then-move like F5-TTS above,
# since it's an occasionally-used feature and shouldn't inflate the packed
# footprint that's materialized on every single GPU call, image generation
# included.
MUSICGEN_MODEL_ID = "facebook/musicgen-small"
from transformers import AutoProcessor as _AutoProcessor, MusicgenForConditionalGeneration

musicgen_processor = _AutoProcessor.from_pretrained(MUSICGEN_MODEL_ID)
musicgen_model = MusicgenForConditionalGeneration.from_pretrained(MUSICGEN_MODEL_ID)
_musicgen_on_gpu = False


# ---------------------------------------------------------------------------
# Tab 1: Image (FLUX.1-dev) + prompt enhancer
# ---------------------------------------------------------------------------

PROMPT_ENHANCER_INSTRUCTION = (
    "Rewrite this into one detailed, vivid image-generation prompt for the FLUX "
    "model. Cover: subject, composition, lens/camera angle, lighting, materials, "
    "mood, and art style. Then on a new line write 'Avoid: ' followed by "
    "comma-separated negatives (blurry, text artifacts, extra limbs, low detail, "
    "watermark, as appropriate — only ones actually relevant to this image). "
    "Output ONLY the prompt and the Avoid line — no preamble, no quotes, no "
    "explanation.\n\nImage idea: "
)


def enhance_prompt(raw_prompt):
    """Calls studio-chat to expand a short prompt. Returns (enhanced_prompt,
    negative_prompt, used_enhancer: bool, note: str)."""
    if not raw_prompt or not raw_prompt.strip():
        return "", "", False, "Empty prompt — nothing to enhance."
    try:
        chat = GradioClient(STUDIO_CHAT_SPACE, hf_token=HF_TOKEN)
        result = chat.predict(
            PROMPT_ENHANCER_INSTRUCTION + raw_prompt,
            "Writing", 300, 0.4,
            api_name="/chat",
        )
        result = result.strip()
        if "\nAvoid:" in result:
            prompt_part, avoid_part = result.split("\nAvoid:", 1)
            return prompt_part.strip(), avoid_part.strip(), True, "Enhanced via studio-chat."
        return result, "", True, "Enhanced via studio-chat (no explicit Avoid line)."
    except Exception as exc:
        return raw_prompt, "", False, f"studio-chat unavailable ({exc}); used raw prompt."


# ---------------------------------------------------------------------------
# Detail Daemon — a faithful port of muerrilla/sd-webui-detail-daemon's
# make_schedule() (MIT-adjacent, code-level technique, no LoRA/extra VRAM).
# Real math ported line-for-line from the original source (verified by
# reading it directly, not guessed from a description), not reimplemented
# from a paraphrase. Applied via callback_on_step_end since FLUX's scheduler
# has use_dynamic_shifting=True — passing a custom `sigmas=` kwarg into the
# pipeline gets re-shifted by the scheduler on top (a real double-shift bug
# avoided by reading the docs before writing any code), so this instead lets
# the pipeline compute its normal, correctly-shifted schedule, then applies
# the multiplier to the already-shifted sigmas in place, once, right after
# step 0 — the same thing the original extension does to the already-final
# per-step sigma at callback time.
# ---------------------------------------------------------------------------


def _detail_daemon_schedule(steps, amount, start=0.2, end=0.8, bias=0.5, exponent=1.0, fade=0.0, smooth=True):
    start = min(start, end)
    mid = start + bias * (end - start)
    multipliers = np.zeros(steps)
    start_idx, mid_idx, end_idx = [int(round(x * (steps - 1))) for x in [start, mid, end]]

    start_values = np.linspace(0, 1, mid_idx - start_idx + 1)
    if smooth:
        start_values = 0.5 * (1 - np.cos(start_values * np.pi))
    start_values = start_values ** exponent
    if start_values.any():
        start_values *= amount

    end_values = np.linspace(1, 0, end_idx - mid_idx + 1)
    if smooth:
        end_values = 0.5 * (1 - np.cos(end_values * np.pi))
    end_values = end_values ** exponent
    if end_values.any():
        end_values *= amount

    multipliers[start_idx:mid_idx + 1] = start_values
    multipliers[mid_idx:end_idx + 1] = end_values
    multipliers *= 1 - fade
    return multipliers


def _make_detail_daemon_callback(amount):
    applied = {"done": False}

    def callback(pipe, step_index, timestep, callback_kwargs):
        if not applied["done"]:
            schedule = _detail_daemon_schedule(len(pipe.scheduler.sigmas) - 1, amount)
            sigmas = pipe.scheduler.sigmas
            for i in range(step_index + 1, len(sigmas) - 1):
                sigmas[i] = sigmas[i] * (1 - schedule[i])
            applied["done"] = True
        return callback_kwargs

    return callback


def _image_duration(model_choice, prompt, negative_prompt, seed, random_seed, num_images, steps, guidance_scale, resolution, detail_boost=False, detail_daemon_amount=0.0):
    # UNMEASURED estimate — the account's daily ZeroGPU quota was exhausted
    # before a real dev generation completed (see README "Known limits").
    # Observed: the platform adds real, unexplained overhead on top of this
    # return value (measured ~52-58s extra on two different calls) — likely
    # GPU materialization cost for this model's packed size, not something
    # this callback controls directly. Re-measure and recalibrate once quota
    # resets; do not trust this as ground truth yet.
    per_step = 3.0
    return min(180, 20 + int(num_images) * int(steps) * per_step)


@spaces.GPU(duration=_image_duration)
def generate_image(model_choice, prompt, negative_prompt, seed, random_seed, num_images, steps, guidance_scale, resolution, detail_boost=False, detail_daemon_amount=0.0):
    pipe = image_pipe
    pipe.set_adapters(["antiblur"], adapter_weights=[1.5 if detail_boost else 0.0])

    width, height = RESOLUTIONS.get(resolution, (1024, 1024))
    num_images = int(num_images)

    seeds = []
    generators = []
    for i in range(num_images):
        s = random.randint(0, 2**31 - 1) if random_seed else int(seed) + i
        seeds.append(s)
        generators.append(torch.Generator("cuda").manual_seed(s))

    use_negative = bool(negative_prompt and negative_prompt.strip())
    kwargs = dict(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=int(steps),
        num_images_per_prompt=num_images,
        generator=generators,
    )
    # dev honors guidance_scale directly (trained guidance-distillation);
    # negative_prompt still needs true_cfg_scale > 1 to take effect.
    kwargs["guidance_scale"] = float(guidance_scale)
    kwargs["true_cfg_scale"] = 1.5 if use_negative else 1.0
    kwargs["negative_prompt"] = negative_prompt if use_negative else None

    detail_daemon_amount = float(detail_daemon_amount)
    if detail_daemon_amount != 0.0:
        kwargs["callback_on_step_end"] = _make_detail_daemon_callback(detail_daemon_amount)
        kwargs["callback_on_step_end_tensor_inputs"] = []

    result = pipe(**kwargs)
    images = result.images + [None] * (4 - len(result.images))
    seed_text = ", ".join(str(s) for s in seeds)
    return images[0], images[1], images[2], images[3], seed_text


def enhance_and_show(raw_prompt):
    enhanced, negative, used, note = enhance_prompt(raw_prompt)
    return enhanced, negative, note


# ---------------------------------------------------------------------------
# Tab 1b: Upscale (Real-ESRGAN, BSD-3-Clause)
# ---------------------------------------------------------------------------


def _upscale_duration(image, scale, restore_faces=False):
    # GFPGAN adds face detection (facexlib/retinaface) + per-face restoration
    # on top of the Real-ESRGAN background upscale — unmeasured margin added
    # here, tightened once a real call's wall-clock time is known.
    base = 20 + int(scale) * 15
    return min(120, base + (40 if restore_faces else 0))


@spaces.GPU(duration=_upscale_duration)
def upscale_image(image, scale, restore_faces=False):
    if image is None:
        raise gr.Error("Upload an image first.")
    # The PyPI package `realesrgan` (installed here) is xinntao's original
    # BasicSR-based release: RealESRGANer(scale, model_path, model=<arch>),
    # not the simpler community RealESRGAN(device).load_weights()/.predict()
    # wrapper this code was first written against — that class doesn't exist
    # in the installed package (real ModuleNotFoundError from a live call).
    import cv2
    from basicsr.archs.rrdbnet_arch import RRDBNet
    from realesrgan import RealESRGANer
    from PIL import Image

    scale = int(scale)
    arch = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=scale)
    model_url = (
        "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
        if scale == 2
        else "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
    )
    upsampler = RealESRGANer(
        scale=scale, model_path=model_url, model=arch,
        tile=0, tile_pad=10, pre_pad=0, half=True, device="cuda",
    )

    img_bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)

    if restore_faces:
        # GFPGAN (Apache-2.0, TencentARC) — already a transitive dependency
        # of realesrgan/basicsr, no new package. Detects + restores faces,
        # then hands the already-constructed Real-ESRGAN upsampler in as its
        # own background upscaler so both run in one pass.
        from gfpgan import GFPGANer

        restorer = GFPGANer(
            model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
            upscale=scale, arch="clean", channel_multiplier=2,
            bg_upsampler=upsampler, device="cuda",
        )
        _, _, output_bgr = restorer.enhance(img_bgr, has_aligned=False, only_center_face=False, paste_back=True)
    else:
        output_bgr, _ = upsampler.enhance(img_bgr, outscale=scale)

    output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(output_rgb)


# ---------------------------------------------------------------------------
# Tab 1c: Remove Background (BiRefNet, MIT — ZhengPeng7/BiRefNet)
# ---------------------------------------------------------------------------


def _bg_removal_duration(image, output_mode):
    return 20


@spaces.GPU(duration=_bg_removal_duration)
def remove_background(image, output_mode):
    if image is None:
        raise gr.Error("Upload an image first.")
    from torchvision import transforms
    from PIL import Image

    src = image.convert("RGB")
    orig_size = src.size

    transform = transforms.Compose([
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    model_dtype = next(birefnet_model.parameters()).dtype
    input_tensor = transform(src).unsqueeze(0).to("cuda", dtype=model_dtype)

    with torch.no_grad():
        mask_logits = birefnet_model(input_tensor)[-1]
        mask = mask_logits.sigmoid().float().cpu()

    mask_img = transforms.ToPILImage()(mask.squeeze(0)).resize(orig_size)

    if output_mode == "Transparent PNG (cutout)":
        result = src.convert("RGBA")
        result.putalpha(mask_img)
        return result
    else:
        # Mask preview — useful to sanity-check the cutout before compositing
        return mask_img.convert("RGB")


# ---------------------------------------------------------------------------
# Tab 2: Transcribe (Whisper-large-v3-turbo) — text + SRT + VTT + word timestamps
# ---------------------------------------------------------------------------


def _fmt_ts(seconds, vtt=False):
    if seconds is None:
        seconds = 0.0
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3_600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    sep = "." if vtt else ","
    return f"{h:02d}:{m:02d}:{s:02d}{sep}{ms:03d}"


def _group_words(words, max_words=9, max_gap=0.6):
    groups, current = [], []
    prev_end = None
    for w in words:
        start, end = w["timestamp"]
        if current and (len(current) >= max_words or (prev_end is not None and start is not None and start - prev_end > max_gap)):
            groups.append(current)
            current = []
        current.append(w)
        prev_end = end
    if current:
        groups.append(current)
    return groups


def _words_to_srt(words):
    groups = _group_words(words)
    lines = []
    for i, g in enumerate(groups, start=1):
        start = g[0]["timestamp"][0] or 0.0
        end = g[-1]["timestamp"][1] or start
        text = "".join(w["text"] for w in g).strip()
        lines.append(f"{i}\n{_fmt_ts(start)} --> {_fmt_ts(end)}\n{text}\n")
    return "\n".join(lines)


def _words_to_vtt(words):
    groups = _group_words(words)
    lines = ["WEBVTT\n"]
    for g in groups:
        start = g[0]["timestamp"][0] or 0.0
        end = g[-1]["timestamp"][1] or start
        text = "".join(w["text"] for w in g).strip()
        lines.append(f"{_fmt_ts(start, vtt=True)} --> {_fmt_ts(end, vtt=True)}\n{text}\n")
    return "\n".join(lines)


def _asr_duration(audio):
    if not audio:
        return 30
    try:
        secs = sf.info(audio).duration
    except Exception:
        secs = 60
    return min(120, 20 + int(secs) // 2)


@spaces.GPU(duration=_asr_duration)
def transcribe(audio):
    if not audio:
        return "", "", "", "[]"
    result = asr_pipe(audio, return_timestamps="word")
    words = result.get("chunks", [])
    text = result.get("text", "").strip()
    srt = _words_to_srt(words)
    vtt = _words_to_vtt(words)
    # json.dumps, never repr() — repr() produces Python-quoted strings, not
    # valid JSON, and silently broke every downstream JSON.parse() caller.
    word_json = json.dumps([
        {"word": w["text"].strip(), "start": w["timestamp"][0], "end": w["timestamp"][1]}
        for w in words
    ])
    return text, srt, vtt, word_json


# ---------------------------------------------------------------------------
# Tab 3: Speak (Kokoro-82M)
# ---------------------------------------------------------------------------


def _tts_duration(text, voice):
    return min(120, 15 + len(text or "") // 15)


@spaces.GPU(duration=_tts_duration)
def speak(text, voice):
    if not text or not text.strip():
        raise gr.Error("Enter some text first.")
    chunks = [audio for _, _, audio in kokoro_pipeline(text, voice=voice)]
    full_audio = np.concatenate(chunks) if chunks else np.zeros(1, dtype=np.float32)
    return (24000, full_audio)


# ---------------------------------------------------------------------------
# Tab 4: Clone (F5-TTS zero-shot voice clone) — Garrett's own voice only
# ---------------------------------------------------------------------------


def _clone_duration(ref_audio, consent, ref_text_override, gen_text):
    return min(120, 30 + len(gen_text or "") // 10)


@spaces.GPU(duration=_clone_duration)
def clone_voice(ref_audio, consent, ref_text_override, gen_text):
    if not consent:
        raise gr.Error(
            "You must check the consent box confirming this is a recording of "
            "your own voice before generating."
        )
    if not ref_audio:
        raise gr.Error("Upload a reference clip of your own voice first.")
    if not gen_text or not gen_text.strip():
        raise gr.Error("Enter text to synthesize.")

    global _f5tts_on_gpu
    if not _f5tts_on_gpu:
        f5tts_model.ema_model.to("cuda")
        f5tts_model.vocoder.to("cuda")
        f5tts_model.device = "cuda"
        _f5tts_on_gpu = True

    ref_text = (ref_text_override or "").strip()
    if not ref_text:
        ref_text = f5tts_model.transcribe(ref_audio)

    wav, sr, _ = f5tts_model.infer(ref_file=ref_audio, ref_text=ref_text, gen_text=gen_text)

    out_dir = tempfile.mkdtemp()
    out_path = os.path.join(out_dir, f"SYNTHESIZED_VOICE_{uuid.uuid4().hex[:8]}.wav")
    sf.write(out_path, wav, sr)
    return out_path


# ---------------------------------------------------------------------------
# Music (MusicGen-small, CC-BY-NC-4.0 weights — non-commercial, personal use)
# ---------------------------------------------------------------------------


def _music_duration(prompt, duration_seconds):
    # UNMEASURED estimate ahead of a real deployed call — MusicGen's own
    # card states ~50 autoregressive steps/second of output, measured on an
    # unspecified (likely stronger-than-A10G) GPU. Conservative multiplier
    # here; recalibrate once a real call's wall-clock time is known, same
    # discipline used for every other UNMEASURED estimate in this Space.
    return min(120, 15 + float(duration_seconds) * 4)


@spaces.GPU(duration=_music_duration)
def generate_music(prompt, duration_seconds):
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the music you want first.")

    global _musicgen_on_gpu
    if not _musicgen_on_gpu:
        musicgen_model.to("cuda")
        _musicgen_on_gpu = True

    inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to("cuda")
    max_new_tokens = min(1500, int(float(duration_seconds) * 50))
    audio_values = musicgen_model.generate(**inputs, max_new_tokens=max_new_tokens)

    sr = musicgen_model.config.audio_encoder.sampling_rate
    out_dir = tempfile.mkdtemp()
    out_path = os.path.join(out_dir, f"music_{uuid.uuid4().hex[:8]}.wav")
    sf.write(out_path, audio_values[0, 0].cpu().numpy(), sr)
    return out_path


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

with gr.Blocks(title="Beexly Studio Media (ZeroGPU)") as demo:
    gr.Markdown(
        "# Beexly Studio Media (ZeroGPU)\n"
        "Personal tool, not part of any Galaxy Sports Edge product, data, or "
        "prediction engine."
    )

    with gr.Tab("Image"):
        gr.Markdown(
            "`FLUX.1-dev` (non-commercial license — personal use only, "
            "**never** for anything shipped on galaxysportsedge.com). "
            "`FLUX.1-schnell` was tried as a second \"fast draft\" model and "
            "reverted — it needs ~24GB more VRAM than this Space has free "
            "once dev is resident. See the README."
        )
        img_model = gr.Dropdown(list(IMAGE_MODELS.keys()), value="Best quality (FLUX.1-dev)", label="Model")
        img_prompt = gr.Textbox(label="Prompt")
        with gr.Row():
            img_enhance_btn = gr.Button("Enhance prompt (via studio-chat)")
        img_enhance_note = gr.Textbox(label="Enhancer status", interactive=False)
        img_negative = gr.Textbox(
            label="Negative prompt (needs true CFG — ~2x slower; auto-filled by enhancer's Avoid line)",
            value="",
        )
        img_resolution = gr.Dropdown(list(RESOLUTIONS.keys()), value="1024x1024 (square)", label="Resolution")
        with gr.Row():
            img_seed = gr.Textbox(label="Seed", value="0")
            img_random_seed = gr.Checkbox(label="Random seed", value=True)
        img_batch = gr.Slider(1, 4, value=1, step=1, label="Batch size")
        img_steps = gr.Slider(20, 50, value=28, step=1, label="Steps (dev only — capped where 1024x1024 stays under 100s; measured post-deploy)")
        img_guidance = gr.Slider(1.0, 7.0, value=3.5, step=0.5, label="Guidance scale (dev only)")
        img_detail_boost = gr.Checkbox(
            label="Detail boost (AntiBlur LoRA, Shakker-Labs — deeper depth of field, sharper background falloff)",
            value=False,
        )
        img_detail_daemon = gr.Slider(
            0.0, 0.3, value=0.0, step=0.02,
            label="Detail Daemon amount (sigma-schedule technique, 0 = off — adjusts mid-sampling noise to add fine texture without changing the prompt)",
        )
        img_btn = gr.Button("Generate", variant="primary")
        img_variations_btn = gr.Button("Variations (4 new seeds, same prompt)")
        with gr.Row():
            img_out1 = gr.Image(label="1")
            img_out2 = gr.Image(label="2")
            img_out3 = gr.Image(label="3")
            img_out4 = gr.Image(label="4")
        img_seeds_used = gr.Textbox(label="Seeds used (reproduce a result by pasting one into Seed and unchecking Random)", interactive=False)

        img_enhance_btn.click(
            enhance_and_show,
            inputs=img_prompt,
            outputs=[img_prompt, img_negative, img_enhance_note],
        )
        img_btn.click(
            generate_image,
            inputs=[img_model, img_prompt, img_negative, img_seed, img_random_seed, img_batch, img_steps, img_guidance, img_resolution, img_detail_boost, img_detail_daemon],
            outputs=[img_out1, img_out2, img_out3, img_out4, img_seeds_used],
            api_name="generate_image",
        )
        img_variations_btn.click(
            generate_image,
            inputs=[img_model, img_prompt, img_negative, img_seed, gr.State(True), gr.State(4), img_steps, img_guidance, img_resolution, img_detail_boost, img_detail_daemon],
            outputs=[img_out1, img_out2, img_out3, img_out4, img_seeds_used],
        )

    with gr.Tab("Upscale"):
        gr.Markdown(
            "`Real-ESRGAN` (BSD-3-Clause) — 2x or 4x upscale of any image. "
            "Optional `GFPGAN` (Apache-2.0) face restoration is for a genuinely "
            "**degraded, blurry, or old** photo of a face — it regenerates "
            "skin texture and color from its own generative prior, not just "
            "sharpens. Tested here on an already-sharp FLUX.1-dev portrait, it "
            "introduced a visible orange/red skin-tone shift — worse, not "
            "better, on a source that wasn't damaged to begin with. **Leave "
            "this off for FLUX output**; it's for restoring a rough source "
            "photo you upload, not for polishing a generated one."
        )
        up_in = gr.Image(type="pil", label="Image to upscale")
        up_scale = gr.Radio(["2", "4"], value="2", label="Scale")
        up_restore_faces = gr.Checkbox(label="Restore faces (GFPGAN — degraded/old photos only, see note above)", value=False)
        up_btn = gr.Button("Upscale")
        up_out = gr.Image(label="Upscaled result")
        up_btn.click(
            upscale_image,
            inputs=[up_in, up_scale, up_restore_faces],
            outputs=up_out,
            api_name="upscale_image",
        )

    with gr.Tab("Remove Background"):
        gr.Markdown(
            "`BiRefNet` (MIT, `ZhengPeng7/BiRefNet`) — subject/background "
            "segmentation. Cuts out the main subject as a transparent PNG, "
            "or shows the raw mask to sanity-check it first."
        )
        bg_in = gr.Image(type="pil", label="Image")
        bg_mode = gr.Radio(
            ["Transparent PNG (cutout)", "Mask only (preview)"],
            value="Transparent PNG (cutout)", label="Output",
        )
        bg_btn = gr.Button("Remove background")
        bg_out = gr.Image(label="Result", image_mode="RGBA")
        bg_btn.click(
            remove_background,
            inputs=[bg_in, bg_mode],
            outputs=bg_out,
            api_name="remove_background",
        )

    with gr.Tab("Transcribe"):
        gr.Markdown(f"Speech-to-text — `{ASR_MODEL_ID}` (MIT). Word-level timestamps, SRT/VTT export.")
        asr_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
        asr_btn = gr.Button("Transcribe")
        asr_text = gr.Textbox(label="Transcript", show_copy_button=True)
        asr_srt = gr.Textbox(label="SRT", lines=6, show_copy_button=True)
        asr_vtt = gr.Textbox(label="VTT", lines=6, show_copy_button=True)
        asr_words = gr.Textbox(label="Word timestamps (JSON)", lines=4, show_copy_button=True)
        asr_btn.click(
            transcribe,
            inputs=asr_input,
            outputs=[asr_text, asr_srt, asr_vtt, asr_words],
            api_name="transcribe",
        )

    with gr.Tab("Speak"):
        gr.Markdown("Text-to-speech — `Kokoro-82M` (Apache-2.0).")
        tts_text = gr.Textbox(label="Text", lines=4)
        tts_voice = gr.Dropdown(TTS_VOICES, value="af_heart", label="Voice")
        tts_btn = gr.Button("Generate speech")
        tts_out = gr.Audio(label="Result")
        tts_btn.click(speak, inputs=[tts_text, tts_voice], outputs=tts_out, api_name="speak")

    with gr.Tab("Clone (your voice only)"):
        gr.Markdown(
            "Zero-shot voice clone — `F5-TTS` (CC-BY-NC-4.0, non-commercial personal "
            "use only). **Upload only a recording of your own voice.** Output "
            "filenames are watermarked `SYNTHESIZED_VOICE_...` so a clip can always "
            "be identified as synthetic."
        )
        clone_ref = gr.Audio(sources=["upload"], type="filepath", label="Reference clip (your own voice)")
        clone_consent = gr.Checkbox(
            label="I confirm this is a recording of my own voice and I consent to it being used for voice synthesis."
        )
        clone_ref_text = gr.Textbox(
            label="Reference transcript (optional — auto-transcribed if left blank)", value=""
        )
        clone_gen_text = gr.Textbox(label="Text to synthesize", lines=4)
        clone_btn = gr.Button("Generate")
        clone_out = gr.Audio(label="Result (synthesized voice)")
        clone_btn.click(
            clone_voice,
            inputs=[clone_ref, clone_consent, clone_ref_text, clone_gen_text],
            outputs=clone_out,
            api_name="clone_voice",
        )

    with gr.Tab("Music"):
        gr.Markdown(
            "`MusicGen-small` (Meta) — text-to-music, e.g. background music "
            "for a video or podcast. **CC-BY-NC-4.0 weights — non-commercial, "
            "personal use only.**"
        )
        music_prompt = gr.Textbox(label="Describe the music", lines=2, placeholder="80s synthwave, driving bassline, upbeat")
        music_duration = gr.Slider(4, 20, value=10, step=1, label="Duration (seconds)")
        music_btn = gr.Button("Generate music")
        music_out = gr.Audio(label="Result")
        music_btn.click(
            generate_music,
            inputs=[music_prompt, music_duration],
            outputs=music_out,
            api_name="generate_music",
        )

if __name__ == "__main__":
    demo.launch()