"""Graphics Enhancer on ZeroGPU: change an image, a video or a piece of music with words.

Image tab (image + text to image): Black Forest Labs' caption model looks at the uploaded picture and writes the full
caption of the edited scene; FLUX.3 Image klein draws it from the picture and that caption.
Video tab (video + text to video): LTX-Video 0.9.8 13B distilled remakes an uploaded clip as you describe,
    keeping its shapes and movement as far as the "how much to change" slider allows.
Music tab (music + text to music): MusicGen Melody writes new music in the style you describe that
    follows the tune of an uploaded piece.

Each model has its own ZeroGPU pack (see zero_packs.py), and each button copies only the packs it needs
onto the GPU, so every job fits the default half-card slice.
"""
import os
import random
import re
import shutil
import subprocess
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 LTXConditionPipeline
from diffusers.utils import export_to_video
from PIL import Image
from transformers import AutoProcessor, MusicgenMelodyForConditionalGeneration

IMAGE_MODEL = "black-forest-labs/flux-3-image-klein"    # Gated, FLUX Kommunity licence: needs the HF_TOKEN secret
CAPTION_MODEL = "black-forest-labs/flux-3-image-klein-pup"   # Apache 2.0: writes FLUX.3's edit captions (sees images)
IMAGE_KLEIN2 = "black-forest-labs/FLUX.2-klein-9B"     # Used if FLUX.3 cannot be loaded (FLUX non-commercial)
IMAGE_FALLBACK = "black-forest-labs/FLUX.2-klein-4B"   # Open (Apache 2.0); used if neither can be loaded
VIDEO_MODEL = "Lightricks/LTX-Video-0.9.8-13B-distilled"
MUSIC_MODEL = "facebook/musicgen-melody"               # Non-commercial licence (CC-BY-NC 4.0)

IMAGE_PIXELS = 1024 * 1024         # Edited pictures keep their shape and are about this many pixels
VIDEO_PIXELS = 768 * 512           # The same for clips
VIDEO_SECONDS = 5                  # Longer clips are cut to their first five seconds
VIDEO_FPS = 16                     # Clips are read at up to this many frames a second
MUSIC_RATE = 32000
MAX_SEED = 2**31 - 1
NEGATIVE_PROMPT = "worst quality, inconsistent motion, blurry, jittery, distorted, watermark, text, deformed"

# 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(...).
# Once a model is sealed, never call .to("cuda") on it again here: its weights are on disk.
print("[load] image model...", flush=True)
# FLUX.3 Image klein (diffusers' modular pipeline). Its caption step runs separately (its own pack, below), so FLUX.3
# and the caption model are never on the GPU together.
from diffusers import Flux2KleinPipeline
try:
    from diffusers import ModularPipeline
    _blocks = ModularPipeline.from_pretrained(IMAGE_MODEL).blocks
    _blocks.sub_blocks.pop("prompt_upsampler")
    image_pipe = _blocks.init_pipeline(IMAGE_MODEL)
    image_pipe.load_components(names=["text_encoder", "tokenizer", "transformer", "vae"], dtype=torch.bfloat16)
    IMAGE_NAME, IMAGE_IS_FLUX3 = IMAGE_MODEL, True
except Exception as e:      # No HF_TOKEN secret, licence not accepted, or diffusers too old: the Space still starts
    print(f"[load] {IMAGE_MODEL} failed ({type(e).__name__}: {str(e)[:200]}); trying {IMAGE_KLEIN2}", flush=True)
    IMAGE_IS_FLUX3 = False
    try:
        image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_KLEIN2, torch_dtype=torch.bfloat16)
        IMAGE_NAME = IMAGE_KLEIN2
    except Exception as e2:
        print(f"[load] {IMAGE_KLEIN2} failed ({type(e2).__name__}: {str(e2)[:200]}); using {IMAGE_FALLBACK}", flush=True)
        image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_FALLBACK, torch_dtype=torch.bfloat16)
        IMAGE_NAME = IMAGE_FALLBACK
print(f"[load] image model in use: {IMAGE_NAME}", flush=True)
for part in (image_pipe.transformer, image_pipe.vae, image_pipe.text_encoder):
    part.to("cuda")
packs.seal("image", IMAGE_NAME)
caption_model = None
if IMAGE_IS_FLUX3:
    print("[load] FLUX.3 caption model...", flush=True)
    try:
        import json
        from huggingface_hub import hf_hub_download
        from transformers import Qwen3_5ForConditionalGeneration
        caption_processor = AutoProcessor.from_pretrained(CAPTION_MODEL)
        caption_model = Qwen3_5ForConditionalGeneration.from_pretrained(CAPTION_MODEL, dtype=torch.bfloat16).eval()
        caption_model.to("cuda")
        with open(hf_hub_download(CAPTION_MODEL, "system_prompts.json"), encoding="utf-8") as f:
            EDIT_RULES = json.load(f)["i2i"]
        packs.seal("caption", CAPTION_MODEL)
    except Exception as e:  # FLUX.3 then gets the instruction as typed (works, but more generic)
        print(f"[load] caption model failed ({type(e).__name__}: {str(e)[:200]})", flush=True)
        caption_model = None
print("[load] LTX...", flush=True)
video_pipe = LTXConditionPipeline.from_pretrained(VIDEO_MODEL, torch_dtype=torch.bfloat16)
video_pipe.vae.enable_tiling()     # Work on 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 Melody...", flush=True)
music_processor = AutoProcessor.from_pretrained(MUSIC_MODEL)
music_model = MusicgenMelodyForConditionalGeneration.from_pretrained(MUSIC_MODEL, torch_dtype=torch.float16).to("cuda")
packs.seal("melody", MUSIC_MODEL)
print("[load] done.", flush=True)

GPU_SIZE = "large" if packs.enabled else "xlarge"


def gpu_seconds(seconds):
    """The GPU time a job asks for. On the half card ZeroGPU counts 1.5 times what is asked, and a job
    that runs past its time is cut off. Without per-model packs every job also waits for all the
    models to be copied, and the full card counts double."""
    return int(seconds) if packs.enabled else int(seconds + 6)


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)


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


def fit(width, height, pixels, multiple):
    """The size closest to the given number of pixels that keeps the shape and suits the model."""
    scale = (pixels / (width * height)) ** 0.5
    return (max(multiple, round(width * scale / multiple) * multiple),
            max(multiple, round(height * scale / multiple) * multiple))


# ---------------------------------------------------------------- Image

def edit_caption(image, instruction):
    """FLUX.3's own caption model looks at the picture and writes the full caption of the edited scene, as Black
    Forest Labs advise (greedy, thinking off). Its pack must already be on the GPU."""
    small = image.copy()
    small.thumbnail((1024, 1024))
    text = ("Reference images follow in order (the first image, the second image, ...).\n"
            f"User request: {instruction}\nOutput the caption.")
    messages = [{"role": "system", "content": EDIT_RULES},
                {"role": "user", "content": [{"type": "image", "image": small}, {"type": "text", "text": text}]}]
    inputs = caption_processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=True,
                                                   return_tensors="pt", enable_thinking=False).to("cuda")
    with torch.no_grad():
        out = caption_model.generate(**inputs, max_new_tokens=1024, do_sample=False)
    return caption_processor.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()


# FLUX.3: the caption (about 8 s) + 16 steps (about 6 s). Estimated: to be checked in the logs.
IMAGE_SECONDS = 26 if IMAGE_IS_FLUX3 else 10


@spaces.GPU(duration=gpu_seconds(IMAGE_SECONDS), size=GPU_SIZE)
def edit_image_gpu(image, instruction, seed):
    with timed("image"):
        width, height = fit(image.width, image.height, IMAGE_PIXELS, 16)
        picture = image.resize((width, height), Image.LANCZOS)
        generator = torch.Generator("cpu").manual_seed(seed)
        if IMAGE_IS_FLUX3:
            caption = instruction
            if caption_model is not None:
                packs.load("caption")
                caption = edit_caption(picture, instruction) or instruction
            packs.load("image")
            return image_pipe(prompt=caption, image=picture, width=width, height=height, generator=generator,
                              output="images")[0], caption
        packs.load("image")
        return image_pipe(image=picture, prompt=instruction, width=width, height=height, num_inference_steps=4,
                          guidance_scale=1.0, generator=generator).images[0], instruction


def edit_image(image, instruction, seed, randomize_seed):
    if image is None:
        raise gr.Error("Upload an image first.")
    if not instruction.strip():
        raise gr.Error("Describe the change you want.")
    seed = pick_seed(seed, randomize_seed)
    picture, caption = edit_image_gpu(image.convert("RGB"), instruction.strip(), seed)
    return picture, seed, caption


# ---------------------------------------------------------------- Video

def read_clip(path):
    """Frames of the first few seconds of a clip, the rate they should be played at, and its length."""
    reader = imageio.get_reader(path)
    source_fps = float(reader.get_meta_data().get("fps") or 24)
    step = max(1, round(source_fps / VIDEO_FPS))
    fps = source_fps / step
    frames = []
    for index, frame in enumerate(reader):
        if index % step == 0:
            frames.append(Image.fromarray(frame).convert("RGB"))
        if len(frames) >= int(VIDEO_SECONDS * fps) + 1:
            break
    reader.close()
    count = (len(frames) - 1) // 8 * 8 + 1          # LTX-Video needs a multiple of 8 frames, plus 1
    return frames[:count], fps


def video_seconds(frames, fps, instruction, change, seed, sound):
    return gpu_seconds(12 + 3.5 * len(frames) / fps)


@spaces.GPU(duration=video_seconds, size=GPU_SIZE)
def edit_video_gpu(frames, fps, instruction, change, seed, sound):
    with timed("video"):
        packs.load("t5", "ltx")
        width, height = fit(frames[0].width, frames[0].height, VIDEO_PIXELS, 32)
        generator = torch.Generator("cpu").manual_seed(seed)
        # The clip itself, in the model's own compressed form, is the starting point
        pixels = video_pipe.video_processor.preprocess_video(frames, height, width).to("cuda", video_pipe.vae.dtype)
        with torch.no_grad():
            latents = video_pipe.vae.encode(pixels).latent_dist.mode()
        latents = video_pipe._normalize_latents(latents, video_pipe.vae.latents_mean, video_pipe.vae.latents_std).float()
        del pixels
        # "change" is how much noise the clip is covered with before it is redrawn: 0.3 keeps almost
        # everything, 0.9 keeps only the rough layout. The first timestep is dropped by denoise_strength.
        result = video_pipe(
            prompt=f"{instruction}. Sharp, crisp, highly detailed, in focus, no blur.",
            negative_prompt=NEGATIVE_PROMPT,
            latents=latents,
            denoise_strength=0.8,
            timesteps=[1000, 1000 * change, 700 * change, 400 * change, 0.03],
            width=width,
            height=height,
            num_frames=len(frames),
            frame_rate=fps,
            guidance_scale=1.0,
            decode_timestep=0.05,
            decode_noise_scale=0.025,
            generator=generator,
        ).frames[0]
        path = f"/tmp/video_{seed}_{random.randint(0, 99999)}.mp4"
        export_to_video(result, path, fps=fps, quality=9)
        return path


def edit_video(video, instruction, change, seed, randomize_seed, keep_sound):
    if not video:
        raise gr.Error("Upload a video first.")
    if not instruction.strip():
        raise gr.Error("Describe how the video should look.")
    frames, fps = read_clip(video)
    if len(frames) < 9:
        raise gr.Error("That video is too short or could not be read. It needs at least half a second.")
    seed = pick_seed(seed, randomize_seed)
    silent = edit_video_gpu(frames, fps, instruction.strip(), float(change), seed, keep_sound)
    if not keep_sound:
        return silent, seed
    # Put the original sound back, if the clip had any
    out = f"/tmp/sound_{random.randint(0, MAX_SEED)}.mp4"
    done = subprocess.run([imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-i", silent, "-i", video, "-map", "0:v:0", "-map", "1:a:0",
                           "-c:v", "copy", "-c:a", "aac", "-shortest", out], capture_output=True)
    return (out if done.returncode == 0 else silent), seed


# ---------------------------------------------------------------- Music

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


def music_seconds(wave, description, seconds):
    return gpu_seconds(5 + seconds * 1.2)


@spaces.GPU(duration=music_seconds, size=GPU_SIZE)
def edit_music_gpu(wave, description, seconds):
    with timed("music"):
        packs.load("melody")
        inputs = music_processor(audio=wave, sampling_rate=MUSIC_RATE, text=[description], padding=True, return_tensors="pt")
        inputs = {k: (v.to("cuda", torch.float16) if v.is_floating_point() else v.to("cuda")) for k, v in inputs.items()}
        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
        out = audio[0, 0].float().cpu().numpy()[:int(seconds * rate)]
        fade = int(0.25 * rate)
        out[:fade] *= np.linspace(0.0, 1.0, fade)
        out[-fade:] *= np.linspace(1.0, 0.0, fade)
        path = f"/tmp/music_{random.randint(0, MAX_SEED)}.wav"
        scipy.io.wavfile.write(path, rate, (np.clip(out, -1, 1) * 32767).astype(np.int16))
        return path


def edit_music(audio, description, seconds):
    if not audio:
        raise gr.Error("Upload some music first.")
    if not description.strip():
        raise gr.Error("Describe the music you want.")
    try:
        wave = load_audio(audio)
    except Exception:
        raise gr.Error("That audio file could not be read.")
    return edit_music_gpu(wave, description.strip(), float(seconds))


# ---------------------------------------------------------------- Showcase

# The owner's hand-picked examples live in the attached bucket (/data/showcase). The bucket is slow
# to read, so each file is copied once to the local disk and shown from there.
SHOWCASE_DIR = os.path.join(os.environ.get("STUDIO_BUCKET", "/data"), "showcase")
SHOWCASE_COPY = "/tmp/showcase"
SHOWCASE_TYPES = (".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm")


def showcase():
    """The showcase files in name order, each with its file name as the caption."""
    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


# ---------------------------------------------------------------- Page

with gr.Blocks(title="Graphics Enhancer") as demo:
    gr.Markdown("# Graphics Enhancer\nChange what you already have, with words: **FLUX.3 Image klein** edits images, "
                "**LTX-Video** remakes video clips, **MusicGen Melody** rewrites music around the same tune.")
    gr.Markdown("Want to make new images, videos, music and narration from nothing but words, or even transcribe audio? Try my other Space, "
                "[Graphics Studio](https://huggingface.co/spaces/vrishv/graphics-studio).")

    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="Your image", type="pil")
                image_text = gr.Textbox(label="What to change", lines=2,
                                        placeholder="Make it a snowy winter evening, with warm light in the windows")
                with gr.Accordion("More options", open=False):
                    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("Change image", variant="primary")
            with gr.Column():
                image_out = gr.Image(label="New image", type="pil")
                image_caption = gr.Textbox(label="Full caption FLUX.3 was given (written from your picture and words)",
                                           lines=4, interactive=False)

    with gr.Tab("Video"):
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(label="Your video (the first 5 seconds are used)", autoplay=False)
                video_text = gr.Textbox(label="How it should look", lines=2,
                                        placeholder="The same scene as a watercolour painting, soft colours")
                video_change = gr.Slider(0.3, 0.9, value=0.6, step=0.05,
                                         label="How much to change (low keeps the clip, high follows the words)")
                video_sound = gr.Checkbox(value=True, label="Keep the original sound")
                with gr.Accordion("More options", open=False):
                    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("Change video", variant="primary")
            with gr.Column():
                video_out = gr.Video(label="New video", autoplay=False, loop=True)

    with gr.Tab("Music"):
        with gr.Row():
            with gr.Column():
                music_in = gr.Audio(label="Your music (the first 30 seconds are used)", sources=["upload"], type="filepath")
                music_text = gr.Textbox(label="The music you want", lines=2,
                                        placeholder="The same tune as a gentle piano piece with soft strings")
                music_length = gr.Slider(5, 28, value=15, step=1, label="Length (seconds)")
                music_button = gr.Button("Change music", variant="primary")
            with gr.Column():
                music_out = gr.Audio(label="New music", type="filepath", autoplay=False)

    with gr.Tab("Showcase"):
        gr.Markdown("Before and after: a few changes made with this Space.")
        showcase_out = gr.Gallery(label="Showcase", columns=2)

    image_button.click(edit_image, [image_in, image_text, image_seed, image_random], [image_out, image_seed, image_caption])
    video_button.click(edit_video, [video_in, video_text, video_change, video_seed, video_random, video_sound],
                       [video_out, video_seed])
    music_button.click(edit_music, [music_in, music_text, music_length], [music_out])
    demo.load(showcase, None, [showcase_out])

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