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

Image tab (image + text to image): FLUX.2 klein edits an uploaded picture as you describe.
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.2-klein-9B"      # Gated, FLUX non-commercial licence: needs the HF_TOKEN secret
IMAGE_FALLBACK = "black-forest-labs/FLUX.2-klein-4B"   # Open (Apache 2.0); used if the 9B cannot 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)
from diffusers import Flux2KleinPipeline
try:
    image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_MODEL, torch_dtype=torch.bfloat16)
    IMAGE_NAME = IMAGE_MODEL
except Exception as e:      # No HF_TOKEN secret, or the licence was not accepted: the Space still starts
    print(f"[load] {IMAGE_MODEL} failed ({type(e).__name__}: {str(e)[:200]}); using {IMAGE_FALLBACK}", flush=True)
    image_pipe = Flux2KleinPipeline.from_pretrained(IMAGE_FALLBACK, torch_dtype=torch.bfloat16)
    IMAGE_NAME = IMAGE_FALLBACK
for part in (image_pipe.transformer, image_pipe.vae, image_pipe.text_encoder):
    part.to("cuda")
packs.seal("image", IMAGE_NAME)
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

@spaces.GPU(duration=gpu_seconds(10), size=GPU_SIZE)
def edit_image_gpu(image, instruction, seed):
    with timed("image"):
        packs.load("image")
        width, height = fit(image.width, image.height, IMAGE_PIXELS, 16)
        return image_pipe(
            image=image.resize((width, height), Image.LANCZOS),
            prompt=instruction,
            width=width,
            height=height,
            num_inference_steps=4,
            guidance_scale=1.0,
            generator=torch.Generator("cpu").manual_seed(seed),
        ).images[0]


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)
    return edit_image_gpu(image.convert("RGB"), instruction.strip(), seed), seed


# ---------------------------------------------------------------- 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.2 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")

    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])
    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])