app.py

import os
import time
import uuid
from pathlib import Path
from typing import Any, Optional

import gradio as gr
from gradio_client import Client, handle_file

# MoviePy 2.x moved VideoFileClip out of moviepy.editor.
try:
    from moviepy import VideoFileClip
except ImportError:
    from moviepy.editor import VideoFileClip


# ── config ────────────────────────────────────────────────────────────
HF_TOKEN = os.environ.get("TOKEN")

OUTPUT_DIR = Path("uploads/output").resolve()
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

client = Client(
    "leonormar009/Gen-Z",
    token=HF_TOKEN,
    download_files=str(OUTPUT_DIR),
)

UTM = (
    "utm_source=hugging_face_space"
    "&utm_medium=banner"
    "&utm_campaign=pro_cta"
)
PRO_URL = f"https://huggingface.co/spaces/leonormar009/Gen-Z?{UTM}"


# ── helpers ───────────────────────────────────────────────────────────
def _get_file_path(value: Any) -> Optional[str]:
    """
    Convert common Gradio file values to a local file path.

    Depending on the Gradio component configuration and Gradio version,
    uploaded files may be returned as a string, pathlib.Path, or dict.
    """
    if value is None:
        return None

    if isinstance(value, (str, os.PathLike)):
        return os.fspath(value)

    if isinstance(value, dict):
        for key in ("path", "filepath", "name"):
            path = value.get(key)
            if path:
                return os.fspath(path)

    # Some Gradio versions return objects with a `path` attribute.
    path = getattr(value, "path", None)
    if path:
        return os.fspath(path)

    return None


def _resize_clip(clip, width: int, height: int):
    """Support both MoviePy 1.x and MoviePy 2.x APIs."""
    if hasattr(clip, "resized"):  # MoviePy 2.x
        return clip.resized(new_size=(width, height))

    # MoviePy 1.x
    return clip.resize(newsize=(width, height))


def _set_clip_fps(clip, fps: int):
    """Support both MoviePy 1.x and MoviePy 2.x APIs."""
    if hasattr(clip, "with_fps"):  # MoviePy 2.x
        return clip.with_fps(fps)

    # MoviePy 1.x
    return clip.set_fps(fps)


def _subclip(clip, start: float, end: float):
    """Support both MoviePy 1.x and MoviePy 2.x APIs."""
    if hasattr(clip, "subclipped"):  # MoviePy 2.x
        return clip.subclipped(start, end)

    # MoviePy 1.x
    return clip.subclip(start, end)


def preprocess_video(
    path: str,
    target_fps: int = 12,
    target_size: int = 800,
    target_len: float = 4,
) -> str:
    """
    Trim, resize, and re-encode a video for the remote API.

    `target_size` is applied to the video's longest edge.
    """
    if not path or not os.path.isfile(path):
        raise ValueError(f"Input video does not exist: {path!r}")

    if target_fps <= 0:
        raise ValueError("target_fps must be greater than zero.")

    if target_size <= 0:
        raise ValueError("target_size must be greater than zero.")

    source_clip = None
    processed_clip = None

    try:
        source_clip = VideoFileClip(path)

        duration = float(source_clip.duration or 0)
        if duration <= 0:
            raise ValueError("The input video has no valid duration.")

        processed_clip = source_clip

        if duration > target_len:
            processed_clip = _subclip(processed_clip, 0, target_len)

        source_width, source_height = map(int, source_clip.size)

        if source_width >= source_height:
            new_width = target_size
            new_height = round(target_size * source_height / source_width)
        else:
            new_height = target_size
            new_width = round(target_size * source_width / source_height)

        # H.264 generally requires even dimensions.
        new_width = max(2, new_width - (new_width % 2))
        new_height = max(2, new_height - (new_height % 2))

        processed_clip = _resize_clip(
            processed_clip,
            width=new_width,
            height=new_height,
        )
        processed_clip = _set_clip_fps(processed_clip, target_fps)

        output_path = OUTPUT_DIR / f"pre_{uuid.uuid4().hex}.mp4"

        # Do not pass `verbose=False`: that argument was removed in MoviePy 2.x.
        processed_clip.write_videofile(
            str(output_path),
            codec="libx264",
            audio_codec="aac",
            fps=target_fps,
            logger=None,
        )

        return str(output_path)

    finally:
        # Closing the clips prevents leaked file handles and locked temp files.
        if processed_clip is not None:
            try:
                processed_clip.close()
            except Exception:
                pass

        if source_clip is not None and source_clip is not processed_clip:
            try:
                source_clip.close()
            except Exception:
                pass


def _find_video_path(value: Any) -> Optional[str]:
    """
    Extract a video path or URL from the different result formats that
    gradio_client may return.
    """
    if value is None:
        return None

    if isinstance(value, (str, os.PathLike)):
        return os.fspath(value)

    if isinstance(value, dict):
        # The API may return {"video": ...}, {"path": ...}, or FileData.
        for key in ("video", "path", "filepath", "url", "name"):
            if key in value:
                result = _find_video_path(value[key])
                if result:
                    return result

        # Handle nested result structures.
        for nested_value in value.values():
            result = _find_video_path(nested_value)
            if result:
                return result

    if isinstance(value, (list, tuple)):
        for item in value:
            result = _find_video_path(item)
            if result:
                return result

    path = getattr(value, "path", None)
    if path:
        return os.fspath(path)

    return None


# ── main generate ─────────────────────────────────────────────────────
def generate(input_image, input_video, gender):
    image_path = _get_file_path(input_image)
    video_path = _get_file_path(input_video)

    if not image_path:
        raise gr.Error("Missing source image.")

    if not video_path:
        raise gr.Error("Missing target video.")

    pre_video = None

    try:
        print("Starting generation...", flush=True)

        pre_video = preprocess_video(video_path)

        job = client.submit(
            input_image=handle_file(image_path),
            input_video=handle_file(pre_video),
            device="cpu",
            selector="many",
            gender=None if gender == "all" else gender,
            race=None,
            order=None,
            api_name="/generate",
        )

        # result() waits for completion and propagates remote errors.
        result = job.result()
        output_video = _find_video_path(result)

        if not output_video:
            raise RuntimeError(
                f"Remote Space returned no usable output video. "
                f"Received: {result!r}"
            )

        print("Generation finished.", flush=True)

        gr.Info(
            "✨ Your preview is ready.<br>"
            "<strong>Get HD</strong> "
            "(4× quality, no watermark, priority) — "
            '<a href="https://www.face-swap.co/'
            '?utm_source=hfspace_deepfakevideo'
            '&utm_medium=info" target="_blank" rel="noopener">'
            "Upgrade on face-swap.co</a>",
            duration=8,
        )

        return output_video

    except gr.Error:
        # Do not wrap an existing Gradio error a second time.
        raise

    except Exception as exc:
        print(f"Generation failed: {exc}", flush=True)
        raise gr.Error(f"Generation failed: {exc}") from exc

    finally:
        # The preprocessed file is only needed while submitting the job.
        if pre_video and os.path.isfile(pre_video):
            try:
                os.remove(pre_video)
            except OSError as exc:
                print(
                    f"Warning: could not remove temporary file "
                    f"{pre_video}: {exc}",
                    flush=True,
                )