init

.gitignore

@@ -0,0 +1,2 @@
+**/__pycache__/
+*.pyc

app.py

@@ -0,0 +1,97 @@
+import gradio as gr
+import spaces
+import torch
+import os
+
+# ==========================================
+# 1. 核心处理函数 (骨架)
+# ==========================================
+@spaces.GPU
+def run_stylization(input_video_path, prompt, noise_scale, shift_gamma, steps, guidance_scale, seed):
+    """
+    这里是实际推理逻辑的占位符。
+    """
+    if not input_video_path:
+        return None
+    
+    print("========== Inference Start ==========")
+    print(f"Video Path: {input_video_path}")
+    print(f"Prompt: {prompt}")
+    print(f"Params: Noise={noise_scale}, Gamma={shift_gamma}, Steps={steps}, CFG={guidance_scale}, Seed={seed}")
+
+    # [Prototype Logic] 直接返回输入视频用于演示
+    return input_video_path
+
+# ==========================================
+# 2. 界面布局 (Gradio Blocks)
+# ==========================================
+
+# 移除 CSS 参数以修复 TypeError
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("# 🎥 ViBT Video Stylization Interface")
+        gr.Markdown("上传视频并设置风格化参数。")
+
+        with gr.Row():
+            # --- 左侧：输入与设置 ---
+            with gr.Column():
+                # 视频输入
+                input_video = gr.Video(label="Source Video", sources=["upload"])
+                
+                # 提示词
+                prompt_input = gr.Textbox(
+                    label="Style Prompt", 
+                    placeholder="e.g., Van Gogh style, cyberpunk city...",
+                    value="Oil painting style, vivid colors"
+                )
+
+                # 高级参数折叠区
+                with gr.Accordion("Advanced Settings", open=True):
+                    with gr.Row():
+                        noise_scale = gr.Slider(
+                            label="Noise Scale", minimum=0.0, maximum=2.0, step=0.1, value=1.0, 
+                            info="Controls how much noise is added."
+                        )
+                        shift_gamma = gr.Slider(
+                            label="Shift Gamma", minimum=1.0, maximum=10.0, step=0.5, value=5.0,
+                            info="Scheduler parameter."
+                        )
+                    
+                    with gr.Row():
+                        num_steps = gr.Slider(
+                            label="Inference Steps", minimum=10, maximum=50, step=1, value=28,
+                            info="More steps = higher quality but slower."
+                        )
+                        guidance_scale = gr.Slider(
+                            label="Guidance Scale (CFG)", minimum=1.0, maximum=20.0, step=0.5, value=1.5,
+                            info="How closely to follow the prompt."
+                        )
+                    
+                    seed = gr.Number(label="Seed", value=42, precision=0)
+
+                # 提交按钮
+                run_btn = gr.Button("Generate Video", variant="primary")
+
+            # --- 右侧：结果输出 ---
+            with gr.Column():
+                output_video = gr.Video(label="Stylized Result", interactive=False)
+
+    # ==========================================
+    # 3. 事件绑定
+    # ==========================================
+    run_btn.click(
+        fn=run_stylization,
+        inputs=[
+            input_video, 
+            prompt_input, 
+            noise_scale, 
+            shift_gamma, 
+            num_steps, 
+            guidance_scale, 
+            seed
+        ],
+        outputs=[output_video]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

vibt/qwen_image.py

@@ -0,0 +1,347 @@
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from PIL import Image
+
+
+from diffusers import QwenImageEditPipeline, QwenImagePipeline
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
+
+from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit import retrieve_latents
+
+
+def encode_vae_image(pipe, image: torch.Tensor, generator: torch.Generator):
+    latent_channels = pipe.vae.config.z_dim if getattr(pipe, "vae", None) else 16
+    image_latents = retrieve_latents(
+        pipe.vae.encode(image), generator=generator, sample_mode="argmax"
+    )
+    latents_mean = (
+        torch.tensor(pipe.vae.config.latents_mean)
+        .view(1, latent_channels, 1, 1, 1)
+        .to(image_latents.device, image_latents.dtype)
+    )
+    latents_std = (
+        torch.tensor(pipe.vae.config.latents_std)
+        .view(1, latent_channels, 1, 1, 1)
+        .to(image_latents.device, image_latents.dtype)
+    )
+    image_latents = (image_latents - latents_mean) / latents_std
+
+    return image_latents
+
+
+@torch.no_grad()
+def encode_image(pipe: QwenImagePipeline, image):
+    width, height = image.size
+    image = pipe.image_processor.preprocess(image, height, width)
+    image = image.to(dtype=pipe.dtype, device=pipe.device).unsqueeze(2)
+    image_latents = encode_vae_image(pipe, image, None)
+
+    image_latent_height, image_latent_width = image_latents.shape[3:]
+    image_latents = pipe._pack_latents(
+        image_latents,
+        1,
+        pipe.transformer.config.in_channels // 4,
+        image_latent_height,
+        image_latent_width,
+    )
+    return image_latents
+
+
+@torch.no_grad()
+def decode_latents_image(pipe: QwenImagePipeline, latents):
+    latents = pipe._unpack_latents(latents, 1024, 1024, pipe.vae_scale_factor)
+    latents = latents.to(pipe.vae.dtype)
+    latents_mean = (
+        torch.tensor(pipe.vae.config.latents_mean)
+        .view(1, pipe.vae.config.z_dim, 1, 1, 1)
+        .to(latents.device, latents.dtype)
+    )
+    latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view(
+        1, pipe.vae.config.z_dim, 1, 1, 1
+    ).to(latents.device, latents.dtype)
+    latents = latents / latents_std + latents_mean
+    image = pipe.vae.decode(latents, return_dict=False)[0][:, :, 0]
+    image = pipe.image_processor.postprocess(image, output_type="pil")
+    return image
+
+
+aspect_ratios = {
+    "1:1": (1328, 1328),
+    "16:9": (1664, 928),
+    "9:16": (928, 1664),
+    "4:3": (1472, 1104),
+    "3:4": (1104, 1472),
+    "3:2": (1584, 1056),
+    "2:3": (1056, 1584),
+}
+
+
+def process_input_img(image):
+    # find the closest aspect ratio
+    w, h = image.size
+    aspect_ratio = w / h
+    closest_ratio = min(
+        aspect_ratios.items(),
+        key=lambda x: abs((x[1][0] / x[1][1]) - aspect_ratio),
+    )
+    target_size = closest_ratio[1]
+    return image.resize(target_size, Image.LANCZOS)
+
+
+@torch.no_grad()
+def qwen_bridge_gen(
+    self: QwenImageEditPipeline,
+    image: Optional[PipelineImageInput] = None,
+    prompt: Union[str, List[str]] = None,
+    negative_prompt: Union[str, List[str]] = None,
+    true_cfg_scale: float = 4.0,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 1.0,
+    num_images_per_prompt: int = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.Tensor] = None,
+    prompt_embeds: Optional[torch.Tensor] = None,
+    prompt_embeds_mask: Optional[torch.Tensor] = None,
+    negative_prompt_embeds: Optional[torch.Tensor] = None,
+    negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+    max_sequence_length: int = 512,
+    # Bridge specific
+    return_trajectory=False,
+):
+    image_size = image[0].size if isinstance(image, list) else image.size
+    calculated_width, calculated_height = image_size
+    height = height or calculated_height
+    width = width or calculated_width
+
+    multiple_of = self.vae_scale_factor * 2
+    width = width // multiple_of * multiple_of
+    height = height // multiple_of * multiple_of
+
+    # 1. Check inputs. Raise error if not correct
+    self.check_inputs(
+        prompt,
+        height,
+        width,
+        negative_prompt=negative_prompt,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        prompt_embeds_mask=prompt_embeds_mask,
+        negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        max_sequence_length=max_sequence_length,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._attention_kwargs = attention_kwargs
+    self._current_timestep = None
+    self._interrupt = False
+
+    # 2. Define call parameters
+    if prompt is not None and isinstance(prompt, str):
+        batch_size = 1
+    elif prompt is not None and isinstance(prompt, list):
+        batch_size = len(prompt)
+    else:
+        batch_size = prompt_embeds.shape[0]
+
+    device = self._execution_device
+    # 3. Preprocess image
+    if image is not None and not (
+        isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels
+    ):
+        image = self.image_processor.resize(image, calculated_height, calculated_width)
+        prompt_image = image
+        image = self.image_processor.preprocess(
+            image, calculated_height, calculated_width
+        )
+        image = image.unsqueeze(2)
+
+    has_neg_prompt = negative_prompt is not None or (
+        negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+    )
+    do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+    prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+        image=prompt_image,
+        prompt=prompt,
+        prompt_embeds=prompt_embeds,
+        prompt_embeds_mask=prompt_embeds_mask,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        max_sequence_length=max_sequence_length,
+    )
+    if do_true_cfg:
+        negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+            image=prompt_image,
+            prompt=negative_prompt,
+            prompt_embeds=negative_prompt_embeds,
+            prompt_embeds_mask=negative_prompt_embeds_mask,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+
+    # 4. Prepare latent variables
+    num_channels_latents = self.transformer.config.in_channels // 4
+    noise, image_latents = self.prepare_latents(
+        image,
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+    latents = image_latents.clone()
+    img_shapes = [
+        [(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]
+    ] * batch_size
+
+    # 5. Prepare timesteps
+    self.scheduler.set_timesteps(num_inference_steps, device=device)
+    timesteps = self.scheduler.timesteps
+    self._num_timesteps = len(timesteps)
+
+    # handle guidance
+    guidance = None
+    txt_seq_lens = (
+        prompt_embeds_mask.sum(dim=1).tolist()
+        if prompt_embeds_mask is not None
+        else None
+    )
+    negative_txt_seq_lens = (
+        negative_prompt_embeds_mask.sum(dim=1).tolist()
+        if negative_prompt_embeds_mask is not None
+        else None
+    )
+
+    trajectory = [latents] if return_trajectory else None
+
+    # 6. Denoising loop
+    with self.progress_bar(total=num_inference_steps) as progress_bar:
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            self._current_timestep = t
+
+            latent_model_input = latents
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            with self.transformer.cache_context("cond"):
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states_mask=prompt_embeds_mask,
+                    encoder_hidden_states=prompt_embeds,
+                    img_shapes=img_shapes,
+                    # img_shapes=[[(1, 64, 64)]],
+                    txt_seq_lens=txt_seq_lens,
+                    attention_kwargs={},
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred[:, : latents.size(1)]
+
+            if do_true_cfg:
+                with self.transformer.cache_context("uncond"):
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        img_shapes=img_shapes,
+                        # img_shapes=[[(1, 64, 64)]],
+                        txt_seq_lens=negative_txt_seq_lens,
+                        attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                comb_pred = neg_noise_pred + true_cfg_scale * (
+                    noise_pred - neg_noise_pred
+                )
+
+                cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                noise_pred = comb_pred * (cond_norm / noise_norm)
+
+            # # step
+            # next_t = timesteps[i + 1] if i < len(timesteps) - 1 else 0
+
+            # sigma_t = t / 1000
+            # sigma_next_t = next_t / 1000
+            # sigma_delta = sigma_next_t - sigma_t
+            # print(
+            #     f"sigma_t: {sigma_t}, sigma_next_t: {sigma_next_t}, sigma_delta: {sigma_delta}"
+            # )
+
+            # noise = torch.randn(
+            #     latents.shape,
+            #     dtype=latents.dtype,
+            #     device=latents.device,
+            #     generator=generator,
+            # )
+            # eta = torch.sqrt(-sigma_delta * sigma_next_t / sigma_t)
+            # # eta = torch.sqrt(-sigma_delta)
+
+            # coef = torch.clip(noise_pred.abs(), 0, 1) if rescale_noise else 1
+            # latents = latents + noise_pred * sigma_delta + sigma * eta * noise * coef
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+            if return_trajectory:
+                trajectory.append(latents)
+
+            # call the callback, if provided
+            progress_bar.update()
+
+    self._current_timestep = None
+    if output_type == "latent":
+        image = latents
+    else:
+
+        def decode_latents(latents, height, width):
+            latents = self._unpack_latents(
+                latents, height, width, self.vae_scale_factor
+            )
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(
+                1, self.vae.config.z_dim, 1, 1, 1
+            ).to(latents.device, latents.dtype)
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+            return image
+
+        image = decode_latents(latents, height, width)
+        trajectory = (
+            [decode_latents(t, height, width)[0] for t in trajectory]
+            if return_trajectory
+            else None
+        )
+
+    # Offload all models
+    self.maybe_free_model_hooks()
+
+    if not return_dict:
+        return (image,)
+
+    if return_trajectory:
+        return QwenImagePipelineOutput(images=image), trajectory
+    else:
+        return QwenImagePipelineOutput(images=image)

vibt/scheduler.py

@@ -0,0 +1,44 @@
+from diffusers.schedulers import UniPCMultistepScheduler
+import torch
+
+
+class ViBTScheduler(UniPCMultistepScheduler):
+    def __init__(self, **kwargs):
+        super().__init__(**{**kwargs, "use_flow_sigmas": True})
+        self.set_parameters()
+
+    def set_parameters(self, noise_scale=1.0, shift_gamma=5.0, seed=None):
+        self.noise_scale = noise_scale
+        self.config.flow_shift = shift_gamma
+        self.generator = (
+            None if seed is None else torch.Generator("cuda").manual_seed(seed)
+        )
+
+    def step(self, model_output, timestep, sample, **kwargs):
+        delta_t = (
+            max(self.timesteps[self.timesteps < timestep]) - timestep
+            if any(self.timesteps < timestep)
+            else -timestep - 1
+        ) / 1000
+
+        current_t = (timestep + 1) / 1000.0
+        eta = (-delta_t * (current_t + delta_t) / current_t) ** 0.5
+
+        noise = torch.randn(
+            sample.shape,
+            generator=self.generator,
+            device=sample.device,
+            dtype=sample.dtype,
+        )
+        latents = sample + delta_t * model_output + eta * self.noise_scale * noise
+
+        return (latents,)
+
+    @classmethod
+    def from_scheduler(
+        cls, scheduler: UniPCMultistepScheduler, noise_scale=1.0, shift_gamma=5.0
+    ):
+        obj = cls.__new__(cls)
+        obj.__dict__ = scheduler.__dict__.copy()
+        obj.set_parameters(noise_scale, shift_gamma)
+        return obj

vibt/wan.py

@@ -0,0 +1,132 @@
+import torch
+import re
+from diffusers import WanPipeline
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+
+
+@torch.no_grad()
+def encode_video(pipe: WanPipeline, video_frames):
+    video_tensor = pipe.video_processor.preprocess_video(video_frames).to(
+        dtype=pipe.dtype, device=pipe.device
+    )
+    posterior = pipe.vae.encode(video_tensor, return_dict=False)[0]
+    z = posterior.mode()
+    latents_mean = (
+        torch.tensor(pipe.vae.config.latents_mean)
+        .view(1, pipe.vae.config.z_dim, 1, 1, 1)
+        .to(z.device, z.dtype)
+    )
+    latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view(
+        1, pipe.vae.config.z_dim, 1, 1, 1
+    ).to(z.device, z.dtype)
+    latents = (z - latents_mean) * latents_std
+    return latents
+
+
+@torch.no_grad()
+def decode_latents(pipe: WanPipeline, latents):
+    latents = latents.to(pipe.vae.dtype)
+    latents_mean = (
+        torch.tensor(pipe.vae.config.latents_mean)
+        .view(1, pipe.vae.config.z_dim, 1, 1, 1)
+        .to(latents.device, latents.dtype)
+    )
+    latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view(
+        1, pipe.vae.config.z_dim, 1, 1, 1
+    ).to(latents.device, latents.dtype)
+    latents = latents / latents_std + latents_mean
+    video = pipe.vae.decode(latents, return_dict=False)[0]
+    video = pipe.video_processor.postprocess_video(video, output_type="np")
+    return video
+
+
+def name_convert(n: str):
+    # blocks.* attention
+    m = re.match(
+        r"blocks\.(\d+)\.(self_attn|cross_attn)\.(q|k|v|o|norm_k|norm_q)\.(weight|bias)",
+        n,
+    )
+    if m:
+        b, kind, comp, suf = m.groups()
+        attn = "attn1" if kind == "self_attn" else "attn2"
+        if comp in ("q", "k", "v"):
+            return f"blocks.{b}.{attn}.to_{comp}.{suf}"
+        if comp == "o":
+            return f"blocks.{b}.{attn}.to_out.0.{suf}"
+        return f"blocks.{b}.{attn}.{comp}.{suf}"
+
+    # blocks.* ffn
+    m = re.match(r"blocks\.(\d+)\.ffn\.(0|2)\.(weight|bias)", n)
+    if m:
+        b, idx, suf = m.groups()
+        if idx == "0":
+            return f"blocks.{b}.ffn.net.0.proj.{suf}"
+        return f"blocks.{b}.ffn.net.2.{suf}"
+
+    # blocks.* norm3/modulation
+    m = re.match(r"blocks\.(\d+)\.norm3\.(weight|bias)", n)
+    if m:
+        b, suf = m.groups()
+        return f"blocks.{b}.norm2.{suf}"
+
+    m = re.match(r"blocks\.(\d+)\.modulation$", n)
+    if m:
+        b = m.group(1)
+        return f"blocks.{b}.scale_shift_table"
+
+    # patch_embedding
+    if n.startswith("patch_embedding."):
+        return n
+
+    # text / time embedding
+    m = re.match(r"text_embedding\.(0|2)\.(weight|bias)", n)
+    if m:
+        idx, suf = m.groups()
+        lin = "linear_1" if idx == "0" else "linear_2"
+        return f"condition_embedder.text_embedder.{lin}.{suf}"
+
+    m = re.match(r"time_embedding\.(0|2)\.(weight|bias)", n)
+    if m:
+        idx, suf = m.groups()
+        lin = "linear_1" if idx == "0" else "linear_2"
+        return f"condition_embedder.time_embedder.{lin}.{suf}"
+
+    m = re.match(r"time_projection\.1\.(weight|bias)", n)
+    if m:
+        suf = m.group(1)
+        return f"condition_embedder.time_proj.{suf}"
+
+    # head
+    if n == "head.head.weight":
+        return "proj_out.weight"
+    if n == "head.head.bias":
+        return "proj_out.bias"
+    if n == "head.modulation":
+        return "scale_shift_table"
+
+    return n
+
+
+def load_vibt_weight(
+    transformer, repo_name="Yuanshi/Bridge", weight_path=None, local_path=None
+):
+    assert (
+        weight_path or local_path
+    ) is not None, "Either weight_path or local_path must be provided."
+
+    tensors = load_file(local_path or hf_hub_download(repo_name, weight_path))
+
+    new_tensors = {}
+
+    for key, value in tensors.items():
+        key = name_convert(key)
+        new_tensors[key] = value
+
+    for name, param in transformer.named_parameters():
+        device, dtype = param.device, param.dtype
+        if name in new_tensors:
+            assert (
+                param.shape == new_tensors[name].shape
+            ), f"{name}: {param.shape} != {new_tensors[name].shape}"
+            param.data = new_tensors[name].to(device=device, dtype=dtype)