Spaces:

hugging-apps
/

vitronedit

Running on Zero

Size GPU durations to measured times; cap frames at 65; better video example

0e2afa4 verified 4 days ago

23 kB

import json
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must be imported before torch / CUDA-touching libs

import random
import re
import tempfile
import time
from types import SimpleNamespace

import einops
import gradio as gr
import imageio
import numpy as np
import torch
from huggingface_hub import snapshot_download
from PIL import Image
from torchvision import transforms

import hyvideo.commons
import hyvideo.pipelines.hunyuan_video_pipeline as hv_pipeline_module
from hyvideo.commons import auto_offload_model
from hyvideo.commons.infer_state import initialize_infer_state
from hyvideo.pipelines.hunyuan_video_pipeline import HunyuanVideo_1_5_Pipeline
from univedit.data.preprocessing import IMAGE_CLIP_FRAMES, resize_and_center_crop_image

# ---------------------------------------------------------------------------
# Weights (downloaded at startup)
# ---------------------------------------------------------------------------
CKPT_DIR = os.path.abspath("ckpts/omniweaving")
BASE_REPO = "tencent/HY-OmniWeaving"
UNIVEDIT_REPO = "SanRiiiii/UniVEdit"
SIGLIP_VISION_CONFIG = {
    "architectures": ["SiglipVisionModel"],
    "attention_dropout": 0.0,
    "hidden_act": "gelu_pytorch_tanh",
    "hidden_size": 1152,
    "image_size": 384,
    "intermediate_size": 4304,
    "layer_norm_eps": 1e-06,
    "model_type": "siglip_vision_model",
    "num_attention_heads": 16,
    "num_channels": 3,
    "num_hidden_layers": 27,
    "patch_size": 14,
}


def _link(src, dst):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    if not os.path.lexists(dst):
        os.symlink(src, dst)


def prepare_weights() -> str:
    # OmniWeaving base: transformer, VAE, scheduler and the fine-tuned MLLM text encoder.
    base = snapshot_download(
        BASE_REPO,
        allow_patterns=[
            "config.json",
            "scheduler/*",
            "transformer/*",
            "vae/*",
            "text_encoder/ckpt/*",
        ],
    )
    _link(os.path.join(base, "config.json"), os.path.join(CKPT_DIR, "config.json"))
    for sub in ("scheduler", "transformer", "vae"):
        _link(os.path.join(base, sub), os.path.join(CKPT_DIR, sub))
    te_ckpt = os.path.join(base, "text_encoder", "ckpt", "text_encoder_model.safetensors")
    _link(te_ckpt, os.path.join(CKPT_DIR, "text_encoder", "ckpt", "text_encoder_model.safetensors"))

    # Qwen2.5-VL-7B-Instruct: config / tokenizer / processor only. The OmniWeaving
    # fine-tuned MLLM checkpoint holds the full weights in HF key format, so it is
    # used directly as the model weights (no need to fetch the 16 GB base model).
    qwen = snapshot_download(
        "Qwen/Qwen2.5-VL-7B-Instruct",
        allow_patterns=["*.json", "*.txt"],
        ignore_patterns=["model.safetensors.index.json"],
    )
    llm_dir = os.path.join(CKPT_DIR, "text_encoder", "llm")
    os.makedirs(llm_dir, exist_ok=True)
    for f in os.listdir(qwen):
        if f != "model.safetensors.index.json":
            _link(os.path.join(qwen, f), os.path.join(llm_dir, f))
    _link(te_ckpt, os.path.join(llm_dir, "model.safetensors"))

    # Glyph-SDXL-v2 byT5 encoder.
    glyph = snapshot_download("multimodalart/glyph-sdxl-v2-byt5-small", allow_patterns=["assets/*", "pytorch_model.bin"])
    glyph_root = os.path.join(CKPT_DIR, "text_encoder", "Glyph-SDXL-v2")
    for f in ("color_idx.json", "multilingual_10-lang_idx.json"):
        _link(os.path.join(glyph, "assets", f), os.path.join(glyph_root, "assets", f))
    _link(os.path.join(glyph, "pytorch_model.bin"), os.path.join(glyph_root, "checkpoints", "byt5_model.pt"))

    # SigLIP vision encoder. OmniWeaving points at FLUX.1-Redux-dev/image_encoder, which is
    # the (gated) bf16 copy of google/siglip-so400m-patch14-384's vision tower; we use the
    # ungated original with the same vision config + preprocessor.
    siglip = snapshot_download(
        "google/siglip-so400m-patch14-384", allow_patterns=["model.safetensors", "preprocessor_config.json"]
    )
    siglip_dir = os.path.join(CKPT_DIR, "vision_encoder", "siglip")
    enc_dir = os.path.join(siglip_dir, "image_encoder")
    os.makedirs(enc_dir, exist_ok=True)
    with open(os.path.join(enc_dir, "config.json"), "w") as f:
        json.dump(SIGLIP_VISION_CONFIG, f)
    _link(os.path.join(siglip, "model.safetensors"), os.path.join(enc_dir, "model.safetensors"))
    _link(
        os.path.join(siglip, "preprocessor_config.json"),
        os.path.join(siglip_dir, "feature_extractor", "preprocessor_config.json"),
    )

    # UniVEdit: LoRA + modality token embeddings.
    univedit = snapshot_download(UNIVEDIT_REPO, allow_patterns=["geo3d/*"])
    return os.path.join(univedit, "geo3d")


UNIVEDIT_DIR = prepare_weights()

# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
# ZeroGPU has no GPU in the main process: report a large card so the pipeline
# picks the non-tiled-small VAE config, and never enables offloading.
_fake_mem = lambda device=None: 96 * 1024**3  # noqa: E731
hyvideo.commons.get_gpu_memory = _fake_mem
hv_pipeline_module.get_gpu_memory = _fake_mem

# Single-process inference: no sequence parallelism. The repo's ParallelDims builds a
# torch DeviceMesh (real CUDA init), which ZeroGPU's main process cannot do, so we
# register an equivalent no-parallelism state directly.
import hyvideo.commons.parallel_states as _ps

setattr(
    _ps,
    "__parallel_dims",
    SimpleNamespace(sp=1, world_size=1, dp_replicate=1, sp_rank=0, sp_group=None, sp_enabled=False, dp_enabled=False),
)
initialize_infer_state(
    SimpleNamespace(
        sage_blocks_range="0-53",
        no_cache_block_id="53",
        use_sageattn=False,
        enable_torch_compile=False,
        enable_cache=False,
        cache_type="deepcache",
        cache_start_step=11,
        cache_end_step=45,
        total_steps=50,
        cache_step_interval=4,
        use_fp8_gemm=False,
        quant_type="fp8-per-token-sgl",
        include_patterns="double_blocks",
    )
)

# The MLLM weights are already loaded from the fine-tuned checkpoint (see above).
HunyuanVideo_1_5_Pipeline._load_text_encoder_ckpt = staticmethod(lambda text_encoder, ckpt_path: None)

pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(
    pretrained_model_name_or_path=CKPT_DIR,
    transformer_dtype=torch.bfloat16,
    device=torch.device("cpu"),
    transformer_init_device=torch.device("cpu"),
    pipeline_config="omniweaving",
)
pipe.apply_infer_optimization(infer_state=None, enable_offloading=False, enable_group_offloading=False)
pipe.load_modality_token_embeddings(UNIVEDIT_DIR, strict=True)
pipe.transformer.load_lora_adapter(
    pretrained_model_name_or_path_or_dict=UNIVEDIT_DIR,
    prefix=None,
    adapter_name="default",
    weight_name="pytorch_lora_weights.safetensors",
    use_safetensors=True,
)
# Merge the LoRA into the base weights (same math, no per-layer PEFT overhead at inference).
pipe.transformer.fuse_lora(lora_scale=1.0)
pipe.transformer.unload_lora()
pipe.to("cuda")
pipe.execution_device = torch.device("cuda")
pipe.noise_init_device = torch.device("cuda")
print("UniVEdit pipeline ready.")

# ---------------------------------------------------------------------------
# Helpers (ported 1:1 from univedit/inference/batch_infer.py)
# ---------------------------------------------------------------------------
IMAGE_EDIT_BUCKETS = ((352, 624), (432, 576), (480, 480), (576, 432), (624, 352))  # (width, height)
VIDEO_EDIT_360P_BUCKETS = {"1:1": (480, 480), "16:9": (352, 624), "9:16": (624, 352)}  # (height, width)
MAX_SEED = 2**31 - 1
DEFAULT_STEPS = 50
DEFAULT_GUIDANCE = 6.0
DEFAULT_SHIFT = 5.0


def replace_reference_placeholder(text: str) -> str:
    text = re.sub(r"\b(?:a|an|the)\s*\[\s*v\s*\*\s*\]", "the reference image", text, flags=re.IGNORECASE)
    text = re.sub(r"\[\s*v\s*\*\s*\]", "the reference image", text, flags=re.IGNORECASE)
    text = re.sub(r"\s+([,.;:!?])", r"\1", text)
    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()


def image_bucket_size(img: Image.Image):
    ratio = img.width / img.height
    w, h = min(IMAGE_EDIT_BUCKETS, key=lambda s: abs((s[0] / s[1]) - ratio))
    return h, w


def video_bucket_size(width: int, height: int):
    ratio = width / height
    _, (h, w) = min(VIDEO_EDIT_360P_BUCKETS.items(), key=lambda it: abs((it[1][1] / it[1][0]) - ratio))
    return h, w


def encode_frames_to_latents(frames: torch.Tensor, *, scale: bool) -> torch.Tensor:
    pixel_values = transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)(frames)
    pixel_values = pixel_values.unsqueeze(0).transpose(1, 2).to(pipe.execution_device)
    with (
        torch.no_grad(),
        torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True),
        auto_offload_model(pipe.vae, pipe.execution_device, enabled=False),
        pipe.vae.memory_efficient_context(),
    ):
        latents = pipe.vae.encode(pixel_values).latent_dist.mode()
        if scale:
            if getattr(pipe.vae.config, "shift_factor", None):
                latents = (latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
            else:
                latents = latents * pipe.vae.config.scaling_factor
    return latents


def to_pil_first_frame(video: torch.Tensor) -> Image.Image:
    frame = video[0, :, 0] if video.ndim == 5 else video[0]
    frame = (frame * 255).clamp(0, 255).to(torch.uint8)
    return Image.fromarray(einops.rearrange(frame, "c h w -> h w c").cpu().numpy())


def resolve_seed(seed, randomize):
    return random.randint(0, MAX_SEED) if randomize else int(seed)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@spaces.GPU(duration=70)  # measured ~60s for 50 steps (+~15%)
def edit_image(
    source_image: Image.Image,
    instruction: str,
    reference_image: Image.Image | None = None,
    steps: int = DEFAULT_STEPS,
    guidance_scale: float = DEFAULT_GUIDANCE,
    flow_shift: float = DEFAULT_SHIFT,
    seed: int = 123,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit an image with a text instruction (optionally guided by a reference image).
    Args:
        source_image: the image to edit.
        instruction: the editing instruction, e.g. "replace the scarf with a hat". With a reference
            image, refer to it as "[V*]" (e.g. "change the dog to a [V*]").
        reference_image: optional reference image (switches to reference-guided editing, ti_i2i).
        steps: number of denoising steps.
        guidance_scale: classifier-free guidance scale.
        flow_shift: timestep shift of the sampling schedule.
        seed: random seed.
        randomize_seed: pick a random seed instead of `seed`.
    Returns:
        The edited image and the seed that was used.
    """
    t0 = time.time()
    if source_image is None:
        raise gr.Error("Please upload an image to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an editing instruction.")
    seed = resolve_seed(seed, randomize_seed)
    source_image = source_image.convert("RGB")
    task = "ti_i2i" if reference_image is not None else "i2i"
    prompt = replace_reference_placeholder(instruction) if task == "ti_i2i" else instruction.strip()

    h, w = image_bucket_size(source_image)
    resized = resize_and_center_crop_image(source_image, h, w)
    clip = transforms.ToTensor()(resized).unsqueeze(0).repeat(IMAGE_CLIP_FRAMES, 1, 1, 1)
    latents = encode_frames_to_latents(clip, scale=True)

    kwargs = dict(
        prompt=prompt,
        negative_prompt="",
        aspect_ratio=f"{w}:{h}",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        flow_shift=float(flow_shift),
        video_length=IMAGE_CLIP_FRAMES,
        seed=seed,
        output_type="pt",
        return_dict=True,
        deepstack=[8, 16, 24],
        setclip=True,
        enable_vae_tile_parallelism=False,  # single process, no torch.distributed
        task_type=task,
        source_image=source_image,
        source_image_latents=latents,
        target_height=h,
        target_width=w,
    )
    if task == "ti_i2i":
        kwargs["reference_image"] = reference_image.convert("RGB")
    t1 = time.time()
    with torch.no_grad():
        out = pipe(**kwargs)
    print(f"[timing] edit_image: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s")
    return to_pil_first_frame(out.videos), seed


def _video_duration(source_video, instruction, reference_image=None, num_frames=33, steps=DEFAULT_STEPS, *args, **kwargs):
    # Calibrated on the live Space: ~1.15 s/step for a 5-frame image clip and ~2.9 s/step
    # for 33 frames at 16:9 (50 steps -> ~148 s); ~10% margin on top.
    lat_frames = (max(5, min(65, int(num_frames))) - 1) // 4 + 1
    tokens = lat_frames * 900
    per_step = 0.55 + 0.000306 * tokens + 1.5e-9 * tokens**2
    return int(5 + 1.1 * int(steps) * per_step)


@spaces.GPU(duration=_video_duration)
def edit_video(
    source_video: str,
    instruction: str,
    reference_image: Image.Image | None = None,
    num_frames: int = 33,
    steps: int = DEFAULT_STEPS,
    guidance_scale: float = DEFAULT_GUIDANCE,
    flow_shift: float = DEFAULT_SHIFT,
    seed: int = 123,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit a short video with a text instruction (optionally guided by a reference image).
    Args:
        source_video: path of the video to edit (the first `num_frames` frames are edited).
        instruction: the editing instruction, e.g. "Replace the brown bear with a panda.".
        reference_image: optional reference image (switches to reference-guided video editing, tiv2v).
        num_frames: number of frames to edit (4n+1, at most 65).
        steps: number of denoising steps.
        guidance_scale: classifier-free guidance scale.
        flow_shift: timestep shift of the sampling schedule.
        seed: random seed.
        randomize_seed: pick a random seed instead of `seed`.
    Returns:
        Path of the edited mp4 video and the seed that was used.
    """
    from decord import VideoReader

    t0 = time.time()
    if not source_video:
        raise gr.Error("Please upload a video to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an editing instruction.")
    seed = resolve_seed(seed, randomize_seed)
    num_frames = int(num_frames)
    num_frames = max(5, min(65, ((num_frames - 1) // 4) * 4 + 1))
    task = "tiv2v" if reference_image is not None else "v2v"

    vr = VideoReader(source_video)
    total = len(vr)
    if total == 0:
        raise gr.Error("The video has no frames.")
    fps = round(vr.get_avg_fps()) or 12
    indices = list(range(min(total, num_frames))) + [total - 1] * max(0, num_frames - total)
    frames = torch.from_numpy(vr.get_batch(indices).asnumpy()).permute(0, 3, 1, 2).contiguous() / 255.0
    del vr

    height, width = frames.shape[-2:]
    th, tw = video_bucket_size(width, height)
    if th / height > tw / width:
        resize_size = th, int(width * th / height)
    else:
        resize_size = int(height * tw / width), tw
    frames = transforms.Compose(
        [transforms.Resize(resize_size, antialias=True), transforms.CenterCrop((th, tw))]
    )(frames)
    latents = encode_frames_to_latents(frames, scale=False)

    kwargs = dict(
        prompt=instruction.strip(),
        negative_prompt="",
        aspect_ratio=f"{tw}:{th}",
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        flow_shift=float(flow_shift),
        video_length=num_frames,
        seed=seed,
        output_type="pt",
        return_dict=True,
        deepstack=[8, 16, 24],
        setclip=True,
        enable_vae_tile_parallelism=False,  # single process, no torch.distributed
        task_type=task,
        source_video=[source_video],
        source_video_latents=latents,
    )
    if task == "tiv2v":
        ref_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
        reference_image.convert("RGB").save(ref_path)
        kwargs["reference_image"] = ref_path
    t1 = time.time()
    with torch.no_grad():
        out = pipe(**kwargs)
    print(f"[timing] edit_video: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s ({num_frames} frames)")

    video = out.videos
    if video.ndim == 5:
        video = video[0]
    vid = (video * 255).clamp(0, 255).to(torch.uint8)
    vid = einops.rearrange(vid, "c f h w -> f h w c").cpu().numpy()
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    imageio.mimwrite(out_path, list(vid), fps=fps, quality=8)
    return out_path, seed


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# VitronEdit (UniVEdit): universal visual editing\n"
            "Edit images and videos with a text instruction, and optionally a reference image. "
            "This is a LoRA plus modality tokens on top of "
            "[HY-OmniWeaving](https://huggingface.co/tencent/HY-OmniWeaving) (HunyuanVideo-1.5).\n\n"
            "[Paper](https://huggingface.co/papers/2610.27049) · "
            "[Code](https://github.com/Vitron-LLM/VitronEdit-Code) · "
            "[Model](https://huggingface.co/SanRiiiii/UniVEdit)"
        )
        with gr.Tabs():
            # ---------------- Image editing ----------------
            with gr.Tab("Image editing"):
                with gr.Row():
                    with gr.Column():
                        img_src = gr.Image(label="Image to edit", type="pil", height=380)
                        img_prompt = gr.Textbox(
                            label="Instruction",
                            placeholder='e.g. "replace the scarf with a hat". With a reference image, call it [V*].',
                            lines=2,
                        )
                        with gr.Accordion("Reference image (optional)", open=False):
                            img_ref = gr.Image(label="Reference image", type="pil", height=240)
                        img_btn = gr.Button("Edit image", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            img_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
                            img_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
                            img_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
                            img_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
                            img_rand = gr.Checkbox(label="Randomize seed", value=False)
                    with gr.Column():
                        img_out = gr.Image(label="Edited image", height=480)
                img_btn.click(
                    edit_image,
                    inputs=[img_src, img_prompt, img_ref, img_steps, img_cfg, img_shift, img_seed, img_rand],
                    outputs=[img_out, img_seed],
                    api_name="edit_image",
                )
                gr.Examples(
                    examples=[
                        ["examples/cat_tabby.jpg", "add a small red bow tie under the cat's chin"],
                        ["examples/living_room_blue_couch.jpg", "replace the blue sofa with a brown leather sofa"],
                        ["examples/rubber_ducks.jpg", "remove the smallest rubber duck on the left"],
                        ["examples/lighthouse_coast.jpg", "Make the style of the image a watercolor painting"],
                    ],
                    inputs=[img_src, img_prompt],
                    outputs=[img_out, img_seed],
                    fn=edit_image,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Image examples",
                )

            # ---------------- Video editing ----------------
            with gr.Tab("Video editing"):
                with gr.Row():
                    with gr.Column():
                        vid_src = gr.Video(label="Video to edit", height=340)
                        vid_prompt = gr.Textbox(
                            label="Instruction",
                            placeholder='e.g. "Replace the brown bear with a panda."',
                            lines=2,
                        )
                        with gr.Accordion("Reference image (optional)", open=False):
                            vid_ref = gr.Image(label="Reference image", type="pil", height=240)
                        vid_btn = gr.Button("Edit video", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            vid_frames = gr.Slider(5, 65, value=33, step=4, label="Frames (4n+1)")
                            vid_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
                            vid_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
                            vid_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
                            vid_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
                            vid_rand = gr.Checkbox(label="Randomize seed", value=False)
                    with gr.Column():
                        vid_out = gr.Video(label="Edited video", autoplay=True, height=420)
                vid_btn.click(
                    edit_video,
                    inputs=[vid_src, vid_prompt, vid_ref, vid_frames, vid_steps, vid_cfg, vid_shift, vid_seed, vid_rand],
                    outputs=[vid_out, vid_seed],
                    api_name="edit_video",
                )
                gr.Examples(
                    examples=[
                        ["examples/cat.mp4", "Replace the cat with a tiger."],
                        ["examples/border_collie_running.mp4", "Replace the dog with a red fox."],
                    ],
                    inputs=[vid_src, vid_prompt],
                    outputs=[vid_out, vid_seed],
                    fn=edit_video,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Video examples",
                )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)

	import json
	import os

	os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

	import spaces # must be imported before torch / CUDA-touching libs

	import random
	import re
	import tempfile
	import time
	from types import SimpleNamespace

	import einops
	import gradio as gr
	import imageio
	import numpy as np
	import torch
	from huggingface_hub import snapshot_download
	from PIL import Image
	from torchvision import transforms

	import hyvideo.commons
	import hyvideo.pipelines.hunyuan_video_pipeline as hv_pipeline_module
	from hyvideo.commons import auto_offload_model
	from hyvideo.commons.infer_state import initialize_infer_state
	from hyvideo.pipelines.hunyuan_video_pipeline import HunyuanVideo_1_5_Pipeline
	from univedit.data.preprocessing import IMAGE_CLIP_FRAMES, resize_and_center_crop_image

	# ---------------------------------------------------------------------------
	# Weights (downloaded at startup)
	# ---------------------------------------------------------------------------
	CKPT_DIR = os.path.abspath("ckpts/omniweaving")
	BASE_REPO = "tencent/HY-OmniWeaving"
	UNIVEDIT_REPO = "SanRiiiii/UniVEdit"
	SIGLIP_VISION_CONFIG = {
	"architectures": ["SiglipVisionModel"],
	"attention_dropout": 0.0,
	"hidden_act": "gelu_pytorch_tanh",
	"hidden_size": 1152,
	"image_size": 384,
	"intermediate_size": 4304,
	"layer_norm_eps": 1e-06,
	"model_type": "siglip_vision_model",
	"num_attention_heads": 16,
	"num_channels": 3,
	"num_hidden_layers": 27,
	"patch_size": 14,
	}


	def _link(src, dst):
	os.makedirs(os.path.dirname(dst), exist_ok=True)
	if not os.path.lexists(dst):
	os.symlink(src, dst)


	def prepare_weights() -> str:
	# OmniWeaving base: transformer, VAE, scheduler and the fine-tuned MLLM text encoder.
	base = snapshot_download(
	BASE_REPO,
	allow_patterns=[
	"config.json",
	"scheduler/*",
	"transformer/*",
	"vae/*",
	"text_encoder/ckpt/*",
	],
	)
	_link(os.path.join(base, "config.json"), os.path.join(CKPT_DIR, "config.json"))
	for sub in ("scheduler", "transformer", "vae"):
	_link(os.path.join(base, sub), os.path.join(CKPT_DIR, sub))
	te_ckpt = os.path.join(base, "text_encoder", "ckpt", "text_encoder_model.safetensors")
	_link(te_ckpt, os.path.join(CKPT_DIR, "text_encoder", "ckpt", "text_encoder_model.safetensors"))

	# Qwen2.5-VL-7B-Instruct: config / tokenizer / processor only. The OmniWeaving
	# fine-tuned MLLM checkpoint holds the full weights in HF key format, so it is
	# used directly as the model weights (no need to fetch the 16 GB base model).
	qwen = snapshot_download(
	"Qwen/Qwen2.5-VL-7B-Instruct",
	allow_patterns=[".json", ".txt"],
	ignore_patterns=["model.safetensors.index.json"],
	)
	llm_dir = os.path.join(CKPT_DIR, "text_encoder", "llm")
	os.makedirs(llm_dir, exist_ok=True)
	for f in os.listdir(qwen):
	if f != "model.safetensors.index.json":
	_link(os.path.join(qwen, f), os.path.join(llm_dir, f))
	_link(te_ckpt, os.path.join(llm_dir, "model.safetensors"))

	# Glyph-SDXL-v2 byT5 encoder.
	glyph = snapshot_download("multimodalart/glyph-sdxl-v2-byt5-small", allow_patterns=["assets/*", "pytorch_model.bin"])
	glyph_root = os.path.join(CKPT_DIR, "text_encoder", "Glyph-SDXL-v2")
	for f in ("color_idx.json", "multilingual_10-lang_idx.json"):
	_link(os.path.join(glyph, "assets", f), os.path.join(glyph_root, "assets", f))
	_link(os.path.join(glyph, "pytorch_model.bin"), os.path.join(glyph_root, "checkpoints", "byt5_model.pt"))

	# SigLIP vision encoder. OmniWeaving points at FLUX.1-Redux-dev/image_encoder, which is
	# the (gated) bf16 copy of google/siglip-so400m-patch14-384's vision tower; we use the
	# ungated original with the same vision config + preprocessor.
	siglip = snapshot_download(
	"google/siglip-so400m-patch14-384", allow_patterns=["model.safetensors", "preprocessor_config.json"]
	)
	siglip_dir = os.path.join(CKPT_DIR, "vision_encoder", "siglip")
	enc_dir = os.path.join(siglip_dir, "image_encoder")
	os.makedirs(enc_dir, exist_ok=True)
	with open(os.path.join(enc_dir, "config.json"), "w") as f:
	json.dump(SIGLIP_VISION_CONFIG, f)
	_link(os.path.join(siglip, "model.safetensors"), os.path.join(enc_dir, "model.safetensors"))
	_link(
	os.path.join(siglip, "preprocessor_config.json"),
	os.path.join(siglip_dir, "feature_extractor", "preprocessor_config.json"),
	)

	# UniVEdit: LoRA + modality token embeddings.
	univedit = snapshot_download(UNIVEDIT_REPO, allow_patterns=["geo3d/*"])
	return os.path.join(univedit, "geo3d")


	UNIVEDIT_DIR = prepare_weights()

	# ---------------------------------------------------------------------------
	# Pipeline
	# ---------------------------------------------------------------------------
	# ZeroGPU has no GPU in the main process: report a large card so the pipeline
	# picks the non-tiled-small VAE config, and never enables offloading.
	_fake_mem = lambda device=None: 96 * 1024**3 # noqa: E731
	hyvideo.commons.get_gpu_memory = _fake_mem
	hv_pipeline_module.get_gpu_memory = _fake_mem

	# Single-process inference: no sequence parallelism. The repo's ParallelDims builds a
	# torch DeviceMesh (real CUDA init), which ZeroGPU's main process cannot do, so we
	# register an equivalent no-parallelism state directly.
	import hyvideo.commons.parallel_states as _ps

	setattr(
	_ps,
	"__parallel_dims",
	SimpleNamespace(sp=1, world_size=1, dp_replicate=1, sp_rank=0, sp_group=None, sp_enabled=False, dp_enabled=False),
	)
	initialize_infer_state(
	SimpleNamespace(
	sage_blocks_range="0-53",
	no_cache_block_id="53",
	use_sageattn=False,
	enable_torch_compile=False,
	enable_cache=False,
	cache_type="deepcache",
	cache_start_step=11,
	cache_end_step=45,
	total_steps=50,
	cache_step_interval=4,
	use_fp8_gemm=False,
	quant_type="fp8-per-token-sgl",
	include_patterns="double_blocks",
	)
	)

	# The MLLM weights are already loaded from the fine-tuned checkpoint (see above).
	HunyuanVideo_1_5_Pipeline._load_text_encoder_ckpt = staticmethod(lambda text_encoder, ckpt_path: None)

	pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(
	pretrained_model_name_or_path=CKPT_DIR,
	transformer_dtype=torch.bfloat16,
	device=torch.device("cpu"),
	transformer_init_device=torch.device("cpu"),
	pipeline_config="omniweaving",
	)
	pipe.apply_infer_optimization(infer_state=None, enable_offloading=False, enable_group_offloading=False)
	pipe.load_modality_token_embeddings(UNIVEDIT_DIR, strict=True)
	pipe.transformer.load_lora_adapter(
	pretrained_model_name_or_path_or_dict=UNIVEDIT_DIR,
	prefix=None,
	adapter_name="default",
	weight_name="pytorch_lora_weights.safetensors",
	use_safetensors=True,
	)
	# Merge the LoRA into the base weights (same math, no per-layer PEFT overhead at inference).
	pipe.transformer.fuse_lora(lora_scale=1.0)
	pipe.transformer.unload_lora()
	pipe.to("cuda")
	pipe.execution_device = torch.device("cuda")
	pipe.noise_init_device = torch.device("cuda")
	print("UniVEdit pipeline ready.")

	# ---------------------------------------------------------------------------
	# Helpers (ported 1:1 from univedit/inference/batch_infer.py)
	# ---------------------------------------------------------------------------
	IMAGE_EDIT_BUCKETS = ((352, 624), (432, 576), (480, 480), (576, 432), (624, 352)) # (width, height)
	VIDEO_EDIT_360P_BUCKETS = {"1:1": (480, 480), "16:9": (352, 624), "9:16": (624, 352)} # (height, width)
	MAX_SEED = 2**31 - 1
	DEFAULT_STEPS = 50
	DEFAULT_GUIDANCE = 6.0
	DEFAULT_SHIFT = 5.0


	def replace_reference_placeholder(text: str) -> str:
	text = re.sub(r"\b(?:a\|an\|the)\s\[\sv\s\\s*\]", "the reference image", text, flags=re.IGNORECASE)
	text = re.sub(r"\[\sv\s\\s\]", "the reference image", text, flags=re.IGNORECASE)
	text = re.sub(r"\s+([,.;:!?])", r"\1", text)
	text = re.sub(r"\s{2,}", " ", text)
	return text.strip()


	def image_bucket_size(img: Image.Image):
	ratio = img.width / img.height
	w, h = min(IMAGE_EDIT_BUCKETS, key=lambda s: abs((s[0] / s[1]) - ratio))
	return h, w


	def video_bucket_size(width: int, height: int):
	ratio = width / height
	_, (h, w) = min(VIDEO_EDIT_360P_BUCKETS.items(), key=lambda it: abs((it[1][1] / it[1][0]) - ratio))
	return h, w


	def encode_frames_to_latents(frames: torch.Tensor, *, scale: bool) -> torch.Tensor:
	pixel_values = transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)(frames)
	pixel_values = pixel_values.unsqueeze(0).transpose(1, 2).to(pipe.execution_device)
	with (
	torch.no_grad(),
	torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True),
	auto_offload_model(pipe.vae, pipe.execution_device, enabled=False),
	pipe.vae.memory_efficient_context(),
	):
	latents = pipe.vae.encode(pixel_values).latent_dist.mode()
	if scale:
	if getattr(pipe.vae.config, "shift_factor", None):
	latents = (latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
	else:
	latents = latents * pipe.vae.config.scaling_factor
	return latents


	def to_pil_first_frame(video: torch.Tensor) -> Image.Image:
	frame = video[0, :, 0] if video.ndim == 5 else video[0]
	frame = (frame * 255).clamp(0, 255).to(torch.uint8)
	return Image.fromarray(einops.rearrange(frame, "c h w -> h w c").cpu().numpy())


	def resolve_seed(seed, randomize):
	return random.randint(0, MAX_SEED) if randomize else int(seed)


	# ---------------------------------------------------------------------------
	# Inference
	# ---------------------------------------------------------------------------
	@spaces.GPU(duration=70) # measured ~60s for 50 steps (+~15%)
	def edit_image(
	source_image: Image.Image,
	instruction: str,
	reference_image: Image.Image \| None = None,
	steps: int = DEFAULT_STEPS,
	guidance_scale: float = DEFAULT_GUIDANCE,
	flow_shift: float = DEFAULT_SHIFT,
	seed: int = 123,
	randomize_seed: bool = False,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Edit an image with a text instruction (optionally guided by a reference image).

	Args:
	source_image: the image to edit.
	instruction: the editing instruction, e.g. "replace the scarf with a hat". With a reference
	image, refer to it as "[V]" (e.g. "change the dog to a [V]").
	reference_image: optional reference image (switches to reference-guided editing, ti_i2i).
	steps: number of denoising steps.
	guidance_scale: classifier-free guidance scale.
	flow_shift: timestep shift of the sampling schedule.
	seed: random seed.
	randomize_seed: pick a random seed instead of `seed`.

	Returns:
	The edited image and the seed that was used.
	"""
	t0 = time.time()
	if source_image is None:
	raise gr.Error("Please upload an image to edit.")
	if not instruction or not instruction.strip():
	raise gr.Error("Please enter an editing instruction.")
	seed = resolve_seed(seed, randomize_seed)
	source_image = source_image.convert("RGB")
	task = "ti_i2i" if reference_image is not None else "i2i"
	prompt = replace_reference_placeholder(instruction) if task == "ti_i2i" else instruction.strip()

	h, w = image_bucket_size(source_image)
	resized = resize_and_center_crop_image(source_image, h, w)
	clip = transforms.ToTensor()(resized).unsqueeze(0).repeat(IMAGE_CLIP_FRAMES, 1, 1, 1)
	latents = encode_frames_to_latents(clip, scale=True)

	kwargs = dict(
	prompt=prompt,
	negative_prompt="",
	aspect_ratio=f"{w}:{h}",
	num_inference_steps=int(steps),
	guidance_scale=float(guidance_scale),
	flow_shift=float(flow_shift),
	video_length=IMAGE_CLIP_FRAMES,
	seed=seed,
	output_type="pt",
	return_dict=True,
	deepstack=[8, 16, 24],
	setclip=True,
	enable_vae_tile_parallelism=False, # single process, no torch.distributed
	task_type=task,
	source_image=source_image,
	source_image_latents=latents,
	target_height=h,
	target_width=w,
	)
	if task == "ti_i2i":
	kwargs["reference_image"] = reference_image.convert("RGB")
	t1 = time.time()
	with torch.no_grad():
	out = pipe(**kwargs)
	print(f"[timing] edit_image: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s")
	return to_pil_first_frame(out.videos), seed


	def _video_duration(source_video, instruction, reference_image=None, num_frames=33, steps=DEFAULT_STEPS, args, *kwargs):
	# Calibrated on the live Space: ~1.15 s/step for a 5-frame image clip and ~2.9 s/step
	# for 33 frames at 16:9 (50 steps -> ~148 s); ~10% margin on top.
	lat_frames = (max(5, min(65, int(num_frames))) - 1) // 4 + 1
	tokens = lat_frames * 900
	per_step = 0.55 + 0.000306 * tokens + 1.5e-9 * tokens**2
	return int(5 + 1.1 * int(steps) * per_step)


	@spaces.GPU(duration=_video_duration)
	def edit_video(
	source_video: str,
	instruction: str,
	reference_image: Image.Image \| None = None,
	num_frames: int = 33,
	steps: int = DEFAULT_STEPS,
	guidance_scale: float = DEFAULT_GUIDANCE,
	flow_shift: float = DEFAULT_SHIFT,
	seed: int = 123,
	randomize_seed: bool = False,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Edit a short video with a text instruction (optionally guided by a reference image).

	Args:
	source_video: path of the video to edit (the first `num_frames` frames are edited).
	instruction: the editing instruction, e.g. "Replace the brown bear with a panda.".
	reference_image: optional reference image (switches to reference-guided video editing, tiv2v).
	num_frames: number of frames to edit (4n+1, at most 65).
	steps: number of denoising steps.
	guidance_scale: classifier-free guidance scale.
	flow_shift: timestep shift of the sampling schedule.
	seed: random seed.
	randomize_seed: pick a random seed instead of `seed`.

	Returns:
	Path of the edited mp4 video and the seed that was used.
	"""
	from decord import VideoReader

	t0 = time.time()
	if not source_video:
	raise gr.Error("Please upload a video to edit.")
	if not instruction or not instruction.strip():
	raise gr.Error("Please enter an editing instruction.")
	seed = resolve_seed(seed, randomize_seed)
	num_frames = int(num_frames)
	num_frames = max(5, min(65, ((num_frames - 1) // 4) * 4 + 1))
	task = "tiv2v" if reference_image is not None else "v2v"

	vr = VideoReader(source_video)
	total = len(vr)
	if total == 0:
	raise gr.Error("The video has no frames.")
	fps = round(vr.get_avg_fps()) or 12
	indices = list(range(min(total, num_frames))) + [total - 1] * max(0, num_frames - total)
	frames = torch.from_numpy(vr.get_batch(indices).asnumpy()).permute(0, 3, 1, 2).contiguous() / 255.0
	del vr

	height, width = frames.shape[-2:]
	th, tw = video_bucket_size(width, height)
	if th / height > tw / width:
	resize_size = th, int(width * th / height)
	else:
	resize_size = int(height * tw / width), tw
	frames = transforms.Compose(
	[transforms.Resize(resize_size, antialias=True), transforms.CenterCrop((th, tw))]
	)(frames)
	latents = encode_frames_to_latents(frames, scale=False)

	kwargs = dict(
	prompt=instruction.strip(),
	negative_prompt="",
	aspect_ratio=f"{tw}:{th}",
	num_inference_steps=int(steps),
	guidance_scale=float(guidance_scale),
	flow_shift=float(flow_shift),
	video_length=num_frames,
	seed=seed,
	output_type="pt",
	return_dict=True,
	deepstack=[8, 16, 24],
	setclip=True,
	enable_vae_tile_parallelism=False, # single process, no torch.distributed
	task_type=task,
	source_video=[source_video],
	source_video_latents=latents,
	)
	if task == "tiv2v":
	ref_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
	reference_image.convert("RGB").save(ref_path)
	kwargs["reference_image"] = ref_path
	t1 = time.time()
	with torch.no_grad():
	out = pipe(**kwargs)
	print(f"[timing] edit_video: prep {t1 - t0:.1f}s, pipe {time.time() - t1:.1f}s ({num_frames} frames)")

	video = out.videos
	if video.ndim == 5:
	video = video[0]
	vid = (video * 255).clamp(0, 255).to(torch.uint8)
	vid = einops.rearrange(vid, "c f h w -> f h w c").cpu().numpy()
	out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
	imageio.mimwrite(out_path, list(vid), fps=fps, quality=8)
	return out_path, seed


	# ---------------------------------------------------------------------------
	# UI
	# ---------------------------------------------------------------------------
	CSS = """
	#col-container { max-width: 1200px; margin: 0 auto; }
	.dark .gradio-container { color: var(--body-text-color); }
	"""

	with gr.Blocks() as demo:
	with gr.Column(elem_id="col-container"):
	gr.Markdown(
	"# VitronEdit (UniVEdit): universal visual editing\n"
	"Edit images and videos with a text instruction, and optionally a reference image. "
	"This is a LoRA plus modality tokens on top of "
	"[HY-OmniWeaving](https://huggingface.co/tencent/HY-OmniWeaving) (HunyuanVideo-1.5).\n\n"
	"[Paper](https://huggingface.co/papers/2610.27049) · "
	"[Code](https://github.com/Vitron-LLM/VitronEdit-Code) · "
	"[Model](https://huggingface.co/SanRiiiii/UniVEdit)"
	)
	with gr.Tabs():
	# ---------------- Image editing ----------------
	with gr.Tab("Image editing"):
	with gr.Row():
	with gr.Column():
	img_src = gr.Image(label="Image to edit", type="pil", height=380)
	img_prompt = gr.Textbox(
	label="Instruction",
	placeholder='e.g. "replace the scarf with a hat". With a reference image, call it [V*].',
	lines=2,
	)
	with gr.Accordion("Reference image (optional)", open=False):
	img_ref = gr.Image(label="Reference image", type="pil", height=240)
	img_btn = gr.Button("Edit image", variant="primary")
	with gr.Accordion("Advanced settings", open=False):
	img_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
	img_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
	img_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
	img_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
	img_rand = gr.Checkbox(label="Randomize seed", value=False)
	with gr.Column():
	img_out = gr.Image(label="Edited image", height=480)
	img_btn.click(
	edit_image,
	inputs=[img_src, img_prompt, img_ref, img_steps, img_cfg, img_shift, img_seed, img_rand],
	outputs=[img_out, img_seed],
	api_name="edit_image",
	)
	gr.Examples(
	examples=[
	["examples/cat_tabby.jpg", "add a small red bow tie under the cat's chin"],
	["examples/living_room_blue_couch.jpg", "replace the blue sofa with a brown leather sofa"],
	["examples/rubber_ducks.jpg", "remove the smallest rubber duck on the left"],
	["examples/lighthouse_coast.jpg", "Make the style of the image a watercolor painting"],
	],
	inputs=[img_src, img_prompt],
	outputs=[img_out, img_seed],
	fn=edit_image,
	cache_examples=True,
	cache_mode="lazy",
	label="Image examples",
	)

	# ---------------- Video editing ----------------
	with gr.Tab("Video editing"):
	with gr.Row():
	with gr.Column():
	vid_src = gr.Video(label="Video to edit", height=340)
	vid_prompt = gr.Textbox(
	label="Instruction",
	placeholder='e.g. "Replace the brown bear with a panda."',
	lines=2,
	)
	with gr.Accordion("Reference image (optional)", open=False):
	vid_ref = gr.Image(label="Reference image", type="pil", height=240)
	vid_btn = gr.Button("Edit video", variant="primary")
	with gr.Accordion("Advanced settings", open=False):
	vid_frames = gr.Slider(5, 65, value=33, step=4, label="Frames (4n+1)")
	vid_steps = gr.Slider(10, 50, value=DEFAULT_STEPS, step=1, label="Steps")
	vid_cfg = gr.Slider(1.0, 10.0, value=DEFAULT_GUIDANCE, step=0.5, label="Guidance scale")
	vid_shift = gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Flow shift")
	vid_seed = gr.Slider(0, MAX_SEED, value=123, step=1, label="Seed")
	vid_rand = gr.Checkbox(label="Randomize seed", value=False)
	with gr.Column():
	vid_out = gr.Video(label="Edited video", autoplay=True, height=420)
	vid_btn.click(
	edit_video,
	inputs=[vid_src, vid_prompt, vid_ref, vid_frames, vid_steps, vid_cfg, vid_shift, vid_seed, vid_rand],
	outputs=[vid_out, vid_seed],
	api_name="edit_video",
	)
	gr.Examples(
	examples=[
	["examples/cat.mp4", "Replace the cat with a tiger."],
	["examples/border_collie_running.mp4", "Replace the dog with a red fox."],
	],
	inputs=[vid_src, vid_prompt],
	outputs=[vid_out, vid_seed],
	fn=edit_video,
	cache_examples=True,
	cache_mode="lazy",
	label="Video examples",
	)

	if __name__ == "__main__":
	demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)