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)