Spaces:

degen69
/

BFS-Video

Runtime error

Upload folder using huggingface_hub

169bbff verified 11 days ago

7.84 kB

"""
BFS — Best Face Swap Video  ·  Hugging Face Space (ZeroGPU)
Swap the identity in any video using the V3 persistent-template technique.
The reference face is placed in a green chroma side-strip that persists across
all frames, giving the model continuous identity conditioning throughout generation.
Model: LTX-2.3 (22B) + BFS head-swap LoRA from degen69/BFS-Best-Face-Swap-Video
Hardware: ZeroGPU xlarge (96 GB)
"""

from __future__ import annotations

# `import spaces` MUST come before torch / diffusers / anything that touches CUDA.
# It patches torch.cuda so the 22B model can load at startup rather than on GPU time.
import spaces

import os
import tempfile
import time

import gradio as gr
import numpy as np
from PIL import Image

from composer import compose_frames, crop_reserved_region
from pipeline import load_pipeline, run_inference
from video_utils import (
    compute_target_size,
    extract_audio,
    frames_for_duration,
    load_video_frames,
    resize_frames,
    save_video,
)

REGION_SIZE = 256
DEFAULT_FPS = 24.0
DEFAULT_DURATION = 5.0
DEFAULT_RESOLUTION = 768

# ---------------------------------------------------------------------------
# Load models once at startup (module scope). `spaces` packs the weights to
# disk; the GPU worker streams them into VRAM on each request.
# ---------------------------------------------------------------------------

print("[bfs] Loading pipeline at startup…", flush=True)
_PIPELINE_STATE = load_pipeline(device="cuda")
print("[bfs] Pipeline ready.", flush=True)


# ---------------------------------------------------------------------------
# Core processing function
# ---------------------------------------------------------------------------

def _estimate_duration(guide_video_path, face_image, prompt, duration, fps, lora_strength, seed, *args, **kwargs) -> int:
    """Scale the GPU reservation with the requested video length."""
    try:
        secs = float(duration)
    except (TypeError, ValueError):
        secs = DEFAULT_DURATION
    # ~60s base + ~45s per second of video, capped at 600 (ZeroGPU max).
    return min(600, int(90 + secs * 45))


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def generate(
    guide_video_path: str,
    face_image: Image.Image,
    prompt: str,
    duration: float,
    fps: float,
    lora_strength: float,
    seed: float,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> tuple[str, str]:
    """
    Full head-swap pipeline:
      1. Load + resize guide video frames
      2. Compose chroma face strip (ReservedRegionFrameComposer)
      3. Run LTX-2.3 diffusion
      4. Crop face strip from output
      5. Mux original audio back in
    Returns (output_video_path, status_message).
    """
    started = time.perf_counter()

    # ---- validate inputs early ----
    if guide_video_path is None:
        return "", "Please upload a guide video."
    if face_image is None:
        return "", "Please upload a reference face image."
    if not prompt.strip():
        return "", "Please enter a text prompt."

    progress(0.05, desc="Loading guide video…")
    frames, source_fps = load_video_frames(guide_video_path)
    if len(frames) == 0:
        return "", "Could not read frames from the guide video."

    # ---- extract audio before we do anything else ----
    audio_tmp = tempfile.mktemp(suffix=".wav")
    has_audio = extract_audio(guide_video_path, audio_tmp)

    # ---- resize frames ----
    progress(0.10, desc="Resizing frames…")
    orig_h, orig_w = frames.shape[1], frames.shape[2]
    target_w, target_h = compute_target_size(orig_w, orig_h, DEFAULT_RESOLUTION)
    frames = resize_frames(frames, target_w, target_h)

    # ---- trim / pad to requested duration ----
    n_frames = frames_for_duration(fps, duration)
    if len(frames) >= n_frames:
        frames = frames[:n_frames]
    else:
        # loop last frame
        pad = np.stack([frames[-1]] * (n_frames - len(frames)))
        frames = np.concatenate([frames, pad], axis=0)

    # ---- compose chroma strip ----
    progress(0.15, desc="Compositing reference face strip…")
    composed = compose_frames(
        frames,
        face_image,
        region_position="left",
        region_size_px=REGION_SIZE,
    )

    # ---- run diffusion ----
    progress(0.20, desc="Running LTX-2.3 diffusion…")
    generated = run_inference(
        _PIPELINE_STATE,
        composed,
        prompt=prompt,
        fps=fps,
        lora_strength=lora_strength,
        seed=int(seed),
        progress_cb=lambda msg: progress(0.20, desc=msg),
    )

    # ---- crop face strip from output ----
    progress(0.90, desc="Cropping reserved region…")
    cropped = crop_reserved_region(
        generated,
        region_position="left",
        region_size_px=REGION_SIZE,
        output_size=(target_w, target_h),
    )

    # ---- save output video with audio ----
    progress(0.95, desc="Encoding output video…")
    out_path = tempfile.mktemp(suffix=".mp4")
    save_video(
        cropped,
        fps=fps,
        output_path=out_path,
        audio_path=audio_tmp if has_audio else None,
        audio_duration=duration,
    )

    elapsed = time.perf_counter() - started
    progress(1.0, desc="Done.")
    return out_path, f"Generation complete in {elapsed:.0f}s."


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
# BFS — Best Face Swap Video
Swap the identity in any video using the **V3 persistent-template** technique.
The reference face is placed in a green chroma side-strip that persists across
all frames, giving the model continuous identity conditioning throughout generation.
**Prompt format:**
```
head_swap:
FACE: Female, fair skin, ~25 years old, long wavy auburn hair, green eyes…
ACTION: A person in a grey hoodie walks toward the camera indoors…
```
"""

with gr.Blocks(title="BFS — Best Face Swap Video") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            guide_video = gr.Video(label="Guide Video", sources=["upload"])
            face_image   = gr.Image(label="Reference Face", type="pil", sources=["upload"])
            prompt       = gr.Textbox(
                label="Text Prompt",
                placeholder="head_swap:\nFACE: ...\nACTION: ...",
                lines=6,
            )

            with gr.Accordion("Parameters", open=False):
                duration       = gr.Slider(1, 15, value=DEFAULT_DURATION, step=0.5,  label="Duration (seconds)")
                fps            = gr.Slider(8, 30,  value=DEFAULT_FPS,      step=1.0,  label="FPS")
                lora_strength  = gr.Slider(0.5, 1.5, value=1.2,            step=0.05, label="Face Swap Strength")
                seed           = gr.Number(value=42, label="Seed", precision=0)

            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            output_video  = gr.Video(label="Result", interactive=False)
            status_text   = gr.Textbox(label="Status", interactive=False)

    run_btn.click(
        fn=generate,
        inputs=[guide_video, face_image, prompt, duration, fps, lora_strength, seed],
        outputs=[output_video, status_text],
        api_name="/generate",
    )

    gr.Markdown("""
---
**Hardware:** ZeroGPU xlarge (96 GB) — billed to the viewer's own HF quota.
**Model:** [degen69/BFS-Best-Face-Swap-Video](https://huggingface.co/degen69/BFS-Best-Face-Swap-Video) · Built on [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3)
**License:** For research and professional VFX use only. You must have explicit consent for any likeness you process.
""")


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

	"""
	BFS — Best Face Swap Video · Hugging Face Space (ZeroGPU)

	Swap the identity in any video using the V3 persistent-template technique.
	The reference face is placed in a green chroma side-strip that persists across
	all frames, giving the model continuous identity conditioning throughout generation.

	Model: LTX-2.3 (22B) + BFS head-swap LoRA from degen69/BFS-Best-Face-Swap-Video
	Hardware: ZeroGPU xlarge (96 GB)
	"""

	from __future__ import annotations

	# `import spaces` MUST come before torch / diffusers / anything that touches CUDA.
	# It patches torch.cuda so the 22B model can load at startup rather than on GPU time.
	import spaces

	import os
	import tempfile
	import time

	import gradio as gr
	import numpy as np
	from PIL import Image

	from composer import compose_frames, crop_reserved_region
	from pipeline import load_pipeline, run_inference
	from video_utils import (
	compute_target_size,
	extract_audio,
	frames_for_duration,
	load_video_frames,
	resize_frames,
	save_video,
	)

	REGION_SIZE = 256
	DEFAULT_FPS = 24.0
	DEFAULT_DURATION = 5.0
	DEFAULT_RESOLUTION = 768

	# ---------------------------------------------------------------------------
	# Load models once at startup (module scope). `spaces` packs the weights to
	# disk; the GPU worker streams them into VRAM on each request.
	# ---------------------------------------------------------------------------

	print("[bfs] Loading pipeline at startup…", flush=True)
	_PIPELINE_STATE = load_pipeline(device="cuda")
	print("[bfs] Pipeline ready.", flush=True)


	# ---------------------------------------------------------------------------
	# Core processing function
	# ---------------------------------------------------------------------------

	def _estimate_duration(guide_video_path, face_image, prompt, duration, fps, lora_strength, seed, args, *kwargs) -> int:
	"""Scale the GPU reservation with the requested video length."""
	try:
	secs = float(duration)
	except (TypeError, ValueError):
	secs = DEFAULT_DURATION
	# ~60s base + ~45s per second of video, capped at 600 (ZeroGPU max).
	return min(600, int(90 + secs * 45))


	@spaces.GPU(duration=_estimate_duration, size="xlarge")
	def generate(
	guide_video_path: str,
	face_image: Image.Image,
	prompt: str,
	duration: float,
	fps: float,
	lora_strength: float,
	seed: float,
	progress: gr.Progress = gr.Progress(track_tqdm=True),
	) -> tuple[str, str]:
	"""
	Full head-swap pipeline:
	1. Load + resize guide video frames
	2. Compose chroma face strip (ReservedRegionFrameComposer)
	3. Run LTX-2.3 diffusion
	4. Crop face strip from output
	5. Mux original audio back in

	Returns (output_video_path, status_message).
	"""
	started = time.perf_counter()

	# ---- validate inputs early ----
	if guide_video_path is None:
	return "", "Please upload a guide video."
	if face_image is None:
	return "", "Please upload a reference face image."
	if not prompt.strip():
	return "", "Please enter a text prompt."

	progress(0.05, desc="Loading guide video…")
	frames, source_fps = load_video_frames(guide_video_path)
	if len(frames) == 0:
	return "", "Could not read frames from the guide video."

	# ---- extract audio before we do anything else ----
	audio_tmp = tempfile.mktemp(suffix=".wav")
	has_audio = extract_audio(guide_video_path, audio_tmp)

	# ---- resize frames ----
	progress(0.10, desc="Resizing frames…")
	orig_h, orig_w = frames.shape[1], frames.shape[2]
	target_w, target_h = compute_target_size(orig_w, orig_h, DEFAULT_RESOLUTION)
	frames = resize_frames(frames, target_w, target_h)

	# ---- trim / pad to requested duration ----
	n_frames = frames_for_duration(fps, duration)
	if len(frames) >= n_frames:
	frames = frames[:n_frames]
	else:
	# loop last frame
	pad = np.stack([frames[-1]] * (n_frames - len(frames)))
	frames = np.concatenate([frames, pad], axis=0)

	# ---- compose chroma strip ----
	progress(0.15, desc="Compositing reference face strip…")
	composed = compose_frames(
	frames,
	face_image,
	region_position="left",
	region_size_px=REGION_SIZE,
	)

	# ---- run diffusion ----
	progress(0.20, desc="Running LTX-2.3 diffusion…")
	generated = run_inference(
	_PIPELINE_STATE,
	composed,
	prompt=prompt,
	fps=fps,
	lora_strength=lora_strength,
	seed=int(seed),
	progress_cb=lambda msg: progress(0.20, desc=msg),
	)

	# ---- crop face strip from output ----
	progress(0.90, desc="Cropping reserved region…")
	cropped = crop_reserved_region(
	generated,
	region_position="left",
	region_size_px=REGION_SIZE,
	output_size=(target_w, target_h),
	)

	# ---- save output video with audio ----
	progress(0.95, desc="Encoding output video…")
	out_path = tempfile.mktemp(suffix=".mp4")
	save_video(
	cropped,
	fps=fps,
	output_path=out_path,
	audio_path=audio_tmp if has_audio else None,
	audio_duration=duration,
	)

	elapsed = time.perf_counter() - started
	progress(1.0, desc="Done.")
	return out_path, f"Generation complete in {elapsed:.0f}s."


	# ---------------------------------------------------------------------------
	# Gradio UI
	# ---------------------------------------------------------------------------

	DESCRIPTION = """
	# BFS — Best Face Swap Video

	Swap the identity in any video using the V3 persistent-template technique.
	The reference face is placed in a green chroma side-strip that persists across
	all frames, giving the model continuous identity conditioning throughout generation.

	Prompt format:
	```
	head_swap:
	FACE: Female, fair skin, ~25 years old, long wavy auburn hair, green eyes…
	ACTION: A person in a grey hoodie walks toward the camera indoors…
	```
	"""

	with gr.Blocks(title="BFS — Best Face Swap Video") as demo:
	gr.Markdown(DESCRIPTION)

	with gr.Row():
	with gr.Column(scale=1):
	guide_video = gr.Video(label="Guide Video", sources=["upload"])
	face_image = gr.Image(label="Reference Face", type="pil", sources=["upload"])
	prompt = gr.Textbox(
	label="Text Prompt",
	placeholder="head_swap:\nFACE: ...\nACTION: ...",
	lines=6,
	)

	with gr.Accordion("Parameters", open=False):
	duration = gr.Slider(1, 15, value=DEFAULT_DURATION, step=0.5, label="Duration (seconds)")
	fps = gr.Slider(8, 30, value=DEFAULT_FPS, step=1.0, label="FPS")
	lora_strength = gr.Slider(0.5, 1.5, value=1.2, step=0.05, label="Face Swap Strength")
	seed = gr.Number(value=42, label="Seed", precision=0)

	run_btn = gr.Button("Generate", variant="primary")

	with gr.Column(scale=1):
	output_video = gr.Video(label="Result", interactive=False)
	status_text = gr.Textbox(label="Status", interactive=False)

	run_btn.click(
	fn=generate,
	inputs=[guide_video, face_image, prompt, duration, fps, lora_strength, seed],
	outputs=[output_video, status_text],
	api_name="/generate",
	)

	gr.Markdown("""
	---
	Hardware: ZeroGPU xlarge (96 GB) — billed to the viewer's own HF quota.
	Model: [degen69/BFS-Best-Face-Swap-Video](https://huggingface.co/degen69/BFS-Best-Face-Swap-Video) · Built on [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3)
	License: For research and professional VFX use only. You must have explicit consent for any likeness you process.
	""")


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