app.py

import os
import subprocess
import tempfile
import traceback

import gradio as gr
import librosa
import numpy as np
import spaces
import torch
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video
from PIL import Image

# ---- КОНФИГ ----------------------------------------------------------
MODEL_ID = "THUDM/CogVideoX-5b-I2V"
NUM_FRAMES = 33          # 33 кадров @ 8fps = ~4 сек клип
FPS = 8
STEPS = 25
GUIDANCE = 6.0
TARGET_W, TARGET_H = 720, 480
DEFAULT_PROMPT = (
    "cinematic smooth camera movement, subtle natural animation, "
    "professional lighting, high quality, detailed"
)

# ---- МОДЕЛЬ ГРУЗИТСЯ СРАЗУ НА СТАРТЕ ---------------------------------
# ZeroGPU поддерживает .to("cuda") на модульном уровне — модель шарится между
# forked worker'ами. Каждый @spaces.GPU вызов тогда — только инференс, без перезаливки
# весов → можно вложиться в duration=90.
print("[boot] loading CogVideoX-5B-I2V pipeline...")
pipe = CogVideoXImageToVideoPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
print("[boot] pipe ready on cuda")


def _letterbox(path: str) -> Image.Image:
    img = Image.open(path).convert("RGB")
    if img.height > img.width:
        tw, th = TARGET_H, TARGET_W
    else:
        tw, th = TARGET_W, TARGET_H
    img2 = img.copy()
    img2.thumbnail((tw, th), Image.LANCZOS)
    canvas = Image.new("RGB", (tw, th), (0, 0, 0))
    canvas.paste(img2, ((tw - img2.width) // 2, (th - img2.height) // 2))
    return canvas


# ---- BEAT DETECTION ---------------------------------------------------------
def _extract_audio(video_path: str) -> str:
    audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    subprocess.run(
        [
            "ffmpeg", "-y", "-i", video_path,
            "-vn", "-ac", "1", "-ar", "22050",
            "-f", "wav", audio_path,
        ],
        check=True, capture_output=True,
    )
    return audio_path


def detect_beats(video_path: str, num_segments: int):
    audio_path = _extract_audio(video_path)
    y, sr = librosa.load(audio_path, sr=None, mono=True)
    duration = float(librosa.get_duration(y=y, sr=sr))
    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
    tempo_val = float(np.atleast_1d(tempo)[0])
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    try:
        os.remove(audio_path)
    except OSError:
        pass

    max_clip = NUM_FRAMES / FPS
    if len(beat_times) < num_segments + 1:
        step = duration / num_segments
        return [max(0.4, min(max_clip, step))] * num_segments, tempo_val, duration

    target_borders = np.linspace(0, duration, num_segments + 1)[1:-1]
    snapped = [0.0]
    for tb in target_borders:
        nearest = beat_times[np.argmin(np.abs(beat_times - tb))]
        snapped.append(float(nearest))
    snapped.append(duration)
    durations = [snapped[i + 1] - snapped[i] for i in range(num_segments)]
    durations = [max(0.4, min(max_clip, d)) for d in durations]
    return durations, tempo_val, duration


# ---- GPU: ОДИН КЛИП ЗА ВЫЗОВ -------------------------------------------
@spaces.GPU(duration=90)
def generate_one_clip(image_path: str, prompt: str, seed: int) -> str:
    image = _letterbox(image_path)
    print(f"[gpu] generating clip (seed={seed}, frames={NUM_FRAMES}, steps={STEPS})...")
    with torch.inference_mode():
        result = pipe(
            prompt=prompt,
            image=image,
            num_videos_per_prompt=1,
            num_inference_steps=STEPS,
            num_frames=NUM_FRAMES,
            guidance_scale=GUIDANCE,
            generator=torch.Generator(device="cuda").manual_seed(seed),
        )
    frames = result.frames[0]
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    export_to_video(frames, out, fps=FPS)
    print(f"[gpu] clip done -> {out}")
    return out


# ---- FFMPEG: TRIM + CONCAT + AUDIO MUX --------------------------------------
def _trim(clip_path: str, dur: float) -> str:
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    subprocess.run(
        [
            "ffmpeg", "-y", "-i", clip_path,
            "-t", f"{dur:.3f}",
            "-vf", "scale=1280:720:force_original_aspect_ratio=decrease,pad=1280:720:(ow-iw)/2:(oh-ih)/2:black,setsar=1",
            "-r", "30",
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
            "-an",
            out,
        ],
        check=True, capture_output=True,
    )
    return out


def _concat_mux(trimmed_clips, ref_video: str) -> str:
    list_txt = tempfile.NamedTemporaryFile(suffix=".txt", delete=False, mode="w").name
    with open(list_txt, "w") as f:
        for p in trimmed_clips:
            f.write(f"file '{p}'\n")

    silent = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    subprocess.run(
        [
            "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", list_txt,
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "20",
            "-an", silent,
        ],
        check=True, capture_output=True,
    )

    final = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    subprocess.run(
        [
            "ffmpeg", "-y",
            "-i", silent, "-i", ref_video,
            "-c:v", "copy",
            "-c:a", "aac", "-b:a", "192k",
            "-map", "0:v:0", "-map", "1:a:0?",
            "-shortest", final,
        ],
        check=False, capture_output=True,
    )
    for p in (list_txt, silent):
        try:
            os.remove(p)
        except OSError:
            pass
    return final


# ---- ОСНОВНОЙ PIPELINE -----------------------------------------------------
def generate(ref_video, img1, img2, img3, img4, img5, img6, img7, style_prompt):
    imgs = [img1, img2, img3, img4, img5, img6, img7]
    if any(x is None for x in imgs):
        raise gr.Error("Нужны все 7 фотографий.")
    if ref_video is None:
        raise gr.Error("Нужно референсное видео для ритма.")

    prompt = (style_prompt or "").strip() or DEFAULT_PROMPT
    log_lines = []

    log_lines.append("[step 1] beat detection...")
    print(log_lines[-1])
    durations, tempo, total_dur = detect_beats(ref_video, num_segments=7)
    log_lines.append(f"[step 1] tempo={tempo:.1f} BPM, ref={total_dur:.2f}s")
    log_lines.append(f"[step 1] segments: {[round(d,2) for d in durations]}")
    print(log_lines[-2]); print(log_lines[-1])

    log_lines.append("[step 2] generating 7 clips (per-clip GPU call, duration=90)...")
    print(log_lines[-1])
    generated = []
    try:
        for i, img in enumerate(imgs, start=1):
            log_lines.append(f"[step 2] clip {i}/7 \u2192 GPU")
            print(log_lines[-1])
            clip = generate_one_clip(img, prompt, 42 + i)
            generated.append(clip)
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Ошибка генерации: {e}")

    log_lines.append("[step 3] trimming to beat segments...")
    print(log_lines[-1])
    trimmed = []
    for i, (clip, dur) in enumerate(zip(generated, durations), start=1):
        log_lines.append(f"[step 3] clip {i}: trim to {dur:.2f}s")
        print(log_lines[-1])
        trimmed.append(_trim(clip, dur))

    log_lines.append("[step 4] concat + audio mux...")
    print(log_lines[-1])
    final = _concat_mux(trimmed, ref_video)

    for p in generated + trimmed:
        try:
            os.remove(p)
        except OSError:
            pass

    log_lines.append("[done] готово")
    print(log_lines[-1])
    return final, "\n".join(log_lines)


# ---- UI ---------------------------------------------------------------------
with gr.Blocks(title="CapCut AI Beat Sync (CogVideoX)") as demo:
    gr.Markdown(
        "## CapCut AI Beat Sync \u2014 self-hosted CogVideoX i2v\n"
        "7 фото + референс → каждая фотка оживает на GPU, клипы режутся по битам, склеиваются с аудио референса."
    )

    ref = gr.Video(label="Референсное видео (бит/ритм)")

    with gr.Row():
        p1 = gr.Image(label="Фото 1", type="filepath")
        p2 = gr.Image(label="Фото 2", type="filepath")
        p3 = gr.Image(label="Фото 3", type="filepath")
        p4 = gr.Image(label="Фото 4", type="filepath")
    with gr.Row():
        p5 = gr.Image(label="Фото 5", type="filepath")
        p6 = gr.Image(label="Фото 6", type="filepath")
        p7 = gr.Image(label="Фото 7", type="filepath")

    prompt = gr.Textbox(
        label="Стиль анимации (опционально)",
        placeholder=DEFAULT_PROMPT,
        lines=2,
    )

    btn = gr.Button("Сгенерировать", variant="primary")

    with gr.Row():
        out_video = gr.Video(label="Результат")
        out_log = gr.Textbox(label="Лог", lines=20, max_lines=40)

    btn.click(
        fn=generate,
        inputs=[ref, p1, p2, p3, p4, p5, p6, p7, prompt],
        outputs=[out_video, out_log],
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)