1 line
160 KiB
JSON
1 line
160 KiB
JSON
{"repo": "aimixer/comfyui-bernini", "files": {"engine/nodes_sampler.py": {"6dd4216a0ba7173a824e5e531a9448e58f623b63": {"t": 1785674302.012098, "n": 159259, "h": "bb6a087e3e58d683", "c": "# Original upstream Copyright (c) 2025 kijai\n# Modified from engine core modules of ComfyUI-WanVideoWrapper\n# Original upstream project: https://github.com/kijai/ComfyUI-WanVideoWrapper\n# Secondary extended Copyright (c) 2026 wuwukasi/wuwukaka\n# Further modified and functionally extended from ComfyUI-WanAnimatePlus\n# Secondary extension project: https://github.com/wuwukaka/ComfyUI-WanAnimatePlus\n# Current modified portions Copyright (c) 2026 AIMixer.\n# - Adapted multi-condition Bernini latent encoding and automatic multi-task recognition logic (T2V/I2V/V2V/RV2V/R2V) for full Bernini pipeline.\n# - Integrated prefix frame multi-reference injection and transition video seamless frame connection mechanism for Bernini video generation and editing scenarios.\n# - Refactored sampler graph-detach protection mechanism to adapt Bernini dual-stage HIGH/LOW denoising workflow and timeline segmented generation.\n# - Optimized context latents and context roles propagation logic for Bernini multi-segment continuous video generation.\n# - Adapted canvas expansion calculation and Uni3C control rendering logic to match Bernini Director node workflow specifications.\n# - Inherited optimized tiled VAE encoding and adaptive VRAM offloading strategies to improve Bernini pipeline memory efficiency.\n# Licensed under the Apache License, Version 2.0\n\nimport os, gc, math, copy\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\nimport inspect\nfrom .wanvideo.modules.model import rope_params\nfrom .custom_linear import remove_lora_from_module, set_lora_params, _replace_linear\nfrom .wanvideo.schedulers import get_scheduler, scheduler_list\nfrom .wanvideo.schedulers.ksampler_bridge import (\n attach_scheduler_timesteps,\n normalize_wanvideo_schedulerv2_inputs,\n resolve_bernini_scheduler,\n resolve_sigma_schedule,\n)\nfrom .gguf.gguf import set_lora_params_gguf\nfrom .utils import(log, print_memory, apply_lora, fourier_filter, optimized_scale, setup_radial_attention,\n compile_model, dict_to_device, tangential_projection, get_raag_guidance, temporal_score_rescaling, offload_transformer, init_blockswap,\n flow_match_add_noise as add_noise)\nfrom .cache_methods.cache_methods import cache_report\nfrom .nodes_model_loading import load_weights\nfrom .enhance_a_video.globals import set_enhance_weight, set_num_frames\nfrom contextlib import nullcontext\nfrom .bernini_guidance import MomentumBuffer, normalized_guidance as _normalized_guidance\nfrom ..bernini.guidance_modes import BERNINI_GUIDANCE_MODES, uses_apg_guidance\n\nfrom comfy import model_management as mm\nfrom comfy.utils import ProgressBar\nfrom comfy.cli_args import args, LatentPreviewMethod\n\nscript_directory = os.path.dirname(os.path.abspath(__file__))\n\ndevice = mm.get_torch_device()\n\n\ndef _latent_noise_mask_to_thw(noise_mask: torch.Tensor) -> torch.Tensor:\n \"\"\"Normalize noise_mask to [T, H, W] (accepts SCAIL-style [1, 1, T, H, W]).\"\"\"\n if noise_mask.ndim == 5:\n return noise_mask.reshape(\n noise_mask.shape[-3], noise_mask.shape[-2], noise_mask.shape[-1]\n )\n if noise_mask.ndim == 4:\n if noise_mask.shape[1] == 1:\n return noise_mask.squeeze(1)\n if noise_mask.shape[0] == 1:\n return noise_mask.squeeze(0)\n return noise_mask\noffload_device = mm.unet_offload_device()\n\nrope_functions = [\"default\", \"comfy\", \"comfy_chunked\"]\n\n\ndef _send_scheduler_sigma_plot(unique_id, sample_scheduler, start_idx, end_idx) -> None:\n try:\n from server import PromptServer\n import io\n import base64\n import matplotlib.pyplot as plt\n except Exception:\n return\n if not unique_id or PromptServer is None:\n return\n try:\n sigmas_np = sample_scheduler.full_sigmas.cpu().numpy()\n if not np.isclose(sigmas_np[-1], 0.0, atol=1e-6):\n sigmas_np = np.append(sigmas_np, 0.0)\n buf = io.BytesIO()\n fig = plt.figure(facecolor='#353535')\n ax = fig.add_subplot(111)\n ax.set_facecolor('#353535')\n x_values = range(0, len(sigmas_np))\n ax.plot(x_values, sigmas_np)\n ax.scatter(x_values, sigmas_np, color='white', s=20, zorder=3)\n for x, y in zip(x_values, sigmas_np):\n show_annotation = len(sigmas_np) <= 10\n is_split_step = (start_idx > 0 and x == start_idx) or (end_idx != -1 and x == end_idx + 1)\n if show_annotation or is_split_step:\n color = 'yellow' if is_split_step else 'orange'\n ax.annotate(\n f\"{y:.3f}\", (x, y), textcoords=\"offset points\",\n xytext=(10, 1), ha='center', color=color, fontsize=12,\n )\n ax.set_xticks(x_values)\n ax.set_title(\"Sigmas\", color='white')\n ax.set_xlabel(\"Step\", color='white')\n ax.set_ylabel(\"Sigma Value\", color='white')\n ax.tick_params(axis='x', colors='white', labelsize=10)\n ax.tick_params(axis='y', colors='white', labelsize=10)\n end_idx_plot = end_idx + 1\n if end_idx_plot != -1 and 0 <= end_idx_plot < len(sigmas_np) - 1:\n ax.axvline(end_idx_plot, color='red', linestyle='--', linewidth=2, label='end_step split')\n if start_idx > 0 and 0 <= start_idx < len(sigmas_np):\n ax.axvline(start_idx, color='green', linestyle='--', linewidth=2, label='start_step split')\n if (end_idx_plot != -1 and 0 <= end_idx_plot < len(sigmas_np)) or (start_idx > 0 and 0 <= start_idx < len(sigmas_np)):\n handles, labels = ax.get_legend_handles_labels()\n if labels:\n ax.legend()\n range_start_idx = start_idx if start_idx > 0 else 0\n range_end_idx = end_idx_plot if end_idx_plot > 0 and end_idx_plot < len(sigmas_np) else len(sigmas_np) - 1\n if range_start_idx < range_end_idx:\n ax.axvspan(range_start_idx, range_end_idx, color='lightblue', alpha=0.1, label='Sampled Range')\n plt.tight_layout()\n plt.savefig(buf, format='png')\n plt.close(fig)\n buf.seek(0)\n img_base64 = base64.b64encode(buf.read()).decode('utf-8')\n buf.close()\n html_img = (\n f\"<img src='data:image/png;base64,{img_base64}' alt='Sigmas Plot' \"\n \"style='max-width:100%; height:100%; overflow:hidden; display:block;'>\"\n )\n PromptServer.instance.send_progress_text(html_img, unique_id)\n except Exception as exc:\n log.error(f\"Failed to send sigmas plot: {exc}\")\n\n\nVAE_STRIDE = (4, 8, 8)\nPATCH_SIZE = (1, 2, 2)\n\n\nclass WanVideoSampler:\n @classmethod\n def INPUT_TYPES(s):\n return {\n \"required\": {\n \"model\": (\"WANVIDEOMODEL\",),\n \"image_embeds\": (\"WANVIDIMAGE_EMBEDS\", ),\n \"steps\": (\"INT\", {\"default\": 30, \"min\": 1}),\n \"cfg\": (\"FLOAT\", {\"default\": 6.0, \"min\": 0.0, \"max\": 30.0, \"step\": 0.01}),\n \"shift\": (\"FLOAT\", {\"default\": 5.0, \"min\": 0.0, \"max\": 1000.0, \"step\": 0.01}),\n \"seed\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 0xffffffffffffffff}),\n \"force_offload\": (\"BOOLEAN\", {\"default\": True, \"tooltip\": \"Moves the model to the offload device after sampling\"}),\n \"scheduler\": (scheduler_list, {\"default\": \"unipc\",}),\n \"riflex_freq_index\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 1000, \"step\": 1, \"tooltip\": \"Frequency index for RIFLEX, disabled when 0, default 6. Allows for new frames to be generated after without looping\"}),\n },\n \"optional\": {\n \"text_embeds\": (\"WANVIDEOTEXTEMBEDS\", ),\n \"samples\": (\"LATENT\", {\"tooltip\": \"init Latents to use for video2video process\"} ),\n \"denoise_strength\": (\"FLOAT\", {\"default\": 1.0, \"min\": 0.0, \"max\": 1.0, \"step\": 0.01}),\n \"feta_args\": (\"FETAARGS\", ),\n \"context_options\": (\"WANVIDCONTEXT\", ),\n \"cache_args\": (\"CACHEARGS\", ),\n \"flowedit_args\": (\"FLOWEDITARGS\", {\"tooltip\": \"FlowEdit support has been deprecated\"}),\n \"batched_cfg\": (\"BOOLEAN\", {\"default\": False, \"tooltip\": \"Batch cond and uncond for faster sampling, possibly faster on some hardware, uses more memory\"}),\n \"slg_args\": (\"SLGARGS\", ),\n \"rope_function\": (rope_functions, {\"default\": \"comfy\", \"tooltip\": \"Comfy's RoPE implementation doesn't use complex numbers and can thus be compiled, that should be a lot faster when using torch.compile. Chunked version has reduced peak VRAM usage when not using torch.compile\"}),\n \"loop_args\": (\"LOOPARGS\", ),\n \"experimental_args\": (\"EXPERIMENTALARGS\", ),\n \"sigmas\": (\"SIGMAS\", ),\n \"unianimate_poses\": (\"UNIANIMATE_POSE\", ),\n \"fantasytalking_embeds\": (\"FANTASYTALKING_EMBEDS\", ),\n \"uni3c_embeds\": (\"UNI3C_EMBEDS\", ),\n \"multitalk_embeds\": (\"MULTITALK_EMBEDS\", ),\n \"freeinit_args\": (\"FREEINITARGS\", ),\n \"start_step\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 10000, \"step\": 1, \"tooltip\": \"Start step for the sampling, 0 means full sampling, otherwise samples only from this step\"}),\n \"end_step\": (\"INT\", {\"default\": -1, \"min\": -1, \"max\": 10000, \"step\": 1, \"tooltip\": \"End step for the sampling, -1 means full sampling, otherwise samples only until this step\"}),\n \"add_noise_to_samples\": (\"BOOLEAN\", {\"default\": False, \"tooltip\": \"Add noise to the samples before sampling, needed for video2video sampling when starting from clean video\"}),\n \"guidance_mode\": (\n BERNINI_GUIDANCE_MODES,\n {\n \"default\": \"none\",\n \"tooltip\": (\n \"Bernini guidance: none = standard CFG; apg / *_apg = Adaptive Projected \"\n \"Guidance; rv2v / v2v / v2v_chain / t2v = official modes (standard CFG in this wrapper).\"\n ),\n },\n ),\n \"apg_eta\": (\"FLOAT\", {\"default\": 0.5, \"min\": 0.0, \"max\": 1.0, \"step\": 0.01, \"tooltip\": \"APG: parallel/orthogonal balance (0=orthogonal only, 1=full)\"}),\n \"apg_momentum\": (\"FLOAT\", {\"default\": -0.5, \"min\": -1.0, \"max\": 1.0, \"step\": 0.01, \"tooltip\": \"APG: EMA momentum for smoothing guidance differences\"}),\n \"apg_norm_threshold\": (\"FLOAT\", {\"default\": 50.0, \"min\": 0.0, \"max\": 1000.0, \"step\": 1.0, \"tooltip\": \"APG: L2 norm clipping threshold (0=disabled)\"}),\n }\n }\n\n RETURN_TYPES = (\"LATENT\", \"LATENT\",)\n RETURN_NAMES = (\"samples\", \"denoised_samples\",)\n FUNCTION = \"process\"\n CATEGORY = \"WanVideoWrapper\"\n\n def process(self, model, image_embeds, shift, steps, cfg, seed, scheduler, riflex_freq_index, text_embeds=None,\n force_offload=True, samples=None, feta_args=None, denoise_strength=1.0, context_options=None,\n cache_args=None, teacache_args=None, flowedit_args=None, batched_cfg=False, slg_args=None, rope_function=\"default\", loop_args=None,\n experimental_args=None, sigmas=None, unianimate_poses=None, fantasytalking_embeds=None, uni3c_embeds=None, multitalk_embeds=None, freeinit_args=None, start_step=0, end_step=-1, add_noise_to_samples=False,\n guidance_mode=\"none\", apg_eta=0.5, apg_momentum=-0.5, apg_norm_threshold=50.0, enable_teacache=None):\n if flowedit_args is not None:\n raise Exception(\"FlowEdit support has been deprecated and removed due to lack of use and code maintainability\")\n _unsupported = (\n (unianimate_poses, \"UniAnimate\"),\n (fantasytalking_embeds, \"FantasyTalking\"),\n (uni3c_embeds, \"Uni3C\"),\n (multitalk_embeds, \"MultiTalk\"),\n )\n for value, label in _unsupported:\n if value is not None:\n raise ValueError(\n f\"{label} is not supported in ComfyUI-Bernini. \"\n \"Use Bernini Context Embeds + Sampler for video workflows.\"\n )\n patcher = model\n model = model.model\n transformer = model.diffusion_model\n\n dtype = model[\"base_dtype\"]\n weight_dtype = model[\"weight_dtype\"]\n fp8_matmul = model[\"fp8_matmul\"]\n gguf_reader = model[\"gguf_reader\"]\n control_lora = model[\"control_lora\"]\n\n vae = image_embeds.get(\"vae\", None)\n tiled_vae = image_embeds.get(\"tiled_vae\", False)\n\n transformer_options = copy.deepcopy(patcher.model_options.get(\"transformer_options\", None))\n merge_loras = transformer_options[\"merge_loras\"]\n\n block_swap_args = transformer_options.get(\"block_swap_args\", None)\n if block_swap_args is not None:\n transformer.use_non_blocking = block_swap_args.get(\"use_non_blocking\", False)\n transformer.blocks_to_swap = block_swap_args.get(\"blocks_to_swap\", 0)\n transformer.vace_blocks_to_swap = block_swap_args.get(\"vace_blocks_to_swap\", 0)\n transformer.prefetch_blocks = block_swap_args.get(\"prefetch_blocks\", 0)\n transformer.block_swap_debug = block_swap_args.get(\"block_swap_debug\", False)\n transformer.offload_img_emb = block_swap_args.get(\"offload_img_emb\", False)\n transformer.offload_txt_emb = block_swap_args.get(\"offload_txt_emb\", False)\n\n is_5b = transformer.out_dim == 48\n vae_upscale_factor = 16 if is_5b else 8\n\n # Load weights\n if transformer.audio_model is not None:\n for block in transformer.blocks:\n if hasattr(block, 'audio_block'):\n block.audio_block = None\n\n if not transformer.patched_linear and patcher.model[\"sd\"] is not None and len(patcher.patches) != 0 and gguf_reader is None:\n transformer = _replace_linear(transformer, dtype, patcher.model[\"sd\"], compile_args=model[\"compile_args\"])\n transformer.patched_linear = True\n if patcher.model[\"sd\"] is not None and gguf_reader is None:\n load_weights(patcher.model.diffusion_model, patcher.model[\"sd\"], weight_dtype, base_dtype=dtype, transformer_load_device=device,\n block_swap_args=block_swap_args, compile_args=model[\"compile_args\"])\n\n if gguf_reader is not None: #handle GGUF\n load_weights(transformer, patcher.model[\"sd\"], base_dtype=dtype, transformer_load_device=device, patcher=patcher, gguf=True,\n reader=gguf_reader, block_swap_args=block_swap_args, compile_args=model[\"compile_args\"])\n set_lora_params_gguf(transformer, patcher.patches)\n transformer.patched_linear = True\n elif len(patcher.patches) != 0: #handle patched linear layers (unmerged loras, fp8 scaled)\n log.info(f\"Using {len(patcher.patches)} LoRA weight patches for WanVideo model\")\n if not merge_loras and fp8_matmul:\n raise NotImplementedError(\"FP8 matmul with unmerged LoRAs is not supported\")\n set_lora_params(transformer, patcher.patches, state_dict=model[\"sd\"], compute_dtype=dtype)\n else:\n remove_lora_from_module(transformer) #clear possible unmerged lora weights\n\n transformer.lora_scheduling_enabled = transformer_options.get(\"lora_scheduling_enabled\", False)\n\n #torch.compile\n if model[\"auto_cpu_offload\"] is False:\n transformer = compile_model(transformer, model[\"compile_args\"])\n\n multitalk_sampling = image_embeds.get(\"multitalk_sampling\", False)\n if multitalk_sampling:\n raise ValueError(\n \"Multitalk is not supported in ComfyUI-Bernini. \"\n \"Use Bernini Context Embeds + Sampler for video editing workflows.\"\n )\n\n if multitalk_sampling and context_options is not None:\n raise Exception(\"context_options are not compatible or necessary with 'WanVideoImageToVideoMultiTalk' node, since it's already an alternative method that creates the video in a loop.\")\n\n if not multitalk_sampling and scheduler == \"multitalk\":\n raise Exception(\"multitalk scheduler is only for multitalk sampling when using ImagetoVideoMultiTalk -node\")\n\n if text_embeds == None:\n text_embeds = {\n \"prompt_embeds\": [],\n \"negative_prompt_embeds\": [],\n }\n else:\n text_embeds = dict_to_device(text_embeds, device)\n\n seed_g = torch.Generator(device=torch.device(\"cpu\"))\n seed_g.manual_seed(seed)\n\n #region Scheduler\n if denoise_strength < 1.0:\n if start_step != 0:\n raise ValueError(\"start_step must be 0 when denoise_strength is used\")\n start_step = steps - int(steps * denoise_strength) - 1\n add_noise_to_samples = True #for now to not break old workflows\n\n sample_scheduler = None\n if isinstance(scheduler, dict):\n sample_scheduler = copy.deepcopy(scheduler[\"sample_scheduler\"])\n timesteps = scheduler[\"timesteps\"]\n stored_sigmas = scheduler.get(\"sigmas\")\n sample_scheduler, timesteps = attach_scheduler_timesteps(\n sample_scheduler,\n timesteps,\n device=device,\n stored_sigmas=stored_sigmas,\n )\n if sample_scheduler.sigmas is not None:\n log.info(\n \"Bernini Sampler loaded scheduler dict: %d step(s), sigmas=%s\",\n len(timesteps),\n [round(float(s), 4) for s in sample_scheduler.sigmas.tolist()],\n )\n start_step = scheduler.get(\"start_step\", start_step)\n elif scheduler != \"multitalk\":\n sample_scheduler, timesteps,_,_ = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, denoise_strength, sigmas=sigmas, log_timesteps=True)\n else:\n timesteps = torch.tensor([1000, 750, 500, 250], device=device)\n\n total_steps = steps\n steps = len(timesteps)\n\n is_pusa = \"pusa\" in sample_scheduler.__class__.__name__.lower()\n\n if scheduler != \"multitalk\":\n scheduler_step_args = {\"generator\": seed_g}\n step_sig = inspect.signature(sample_scheduler.step)\n for arg in list(scheduler_step_args.keys()):\n if arg not in step_sig.parameters:\n scheduler_step_args.pop(arg)\n\n # Ovi\n if transformer.audio_model is not None: # temporary workaround (...nothing more permanent)\n for i, block in enumerate(transformer.blocks):\n block.audio_block = transformer.audio_model.blocks[i]\n sample_scheduler_ovi = copy.deepcopy(sample_scheduler)\n rope_function = \"default\" # comfy rope not implemented for ovi model yet\n ovi_negative_text_embeds = text_embeds.get(\"ovi_negative_prompt_embeds\", None)\n ovi_audio_cfg = text_embeds.get(\"ovi_audio_cfg\", None)\n if ovi_audio_cfg is not None:\n if not isinstance(ovi_audio_cfg, list):\n ovi_audio_cfg = [ovi_audio_cfg] * (steps + 1)\n\n if isinstance(cfg, list):\n if steps < len(cfg):\n log.info(f\"Received {len(cfg)} cfg values, but only {steps} steps. Slicing cfg list to match steps.\")\n cfg = cfg[:steps]\n elif steps > len(cfg):\n log.info(f\"Received only {len(cfg)} cfg values, but {steps} steps. Extending cfg list to match steps.\")\n cfg.extend([cfg[-1]] * (steps - len(cfg)))\n log.info(f\"Using per-step cfg list: {cfg}\")\n else:\n cfg = [cfg] * (steps + 1)\n\n control_latents = control_camera_latents = clip_fea = clip_fea_neg = end_image = recammaster = camera_embed = unianim_data = mocha_embeds = image_cond_neg =None\n vace_data = vace_context = vace_scale = None\n fun_or_fl2v_model = drop_last = False\n phantom_latents = fun_ref_image = ATI_tracks = None\n add_cond = attn_cond = attn_cond_neg = noise_pred_flipped = None\n humo_audio = humo_audio_neg = None\n has_ref = image_embeds.get(\"has_ref\", False)\n\n #I2V\n story_mem_latents = image_embeds.get(\"story_mem_latents\", None)\n image_cond = image_embeds.get(\"image_embeds\", None)\n image_cond_mask = None\n if image_cond is not None:\n if transformer.in_dim == 16:\n raise ValueError(\"T2V (text to video) model detected, encoded images only work with I2V (Image to video) models\")\n elif transformer.in_dim not in [48, 32]: # fun 2.1 models don't use the mask\n image_cond_mask = image_embeds.get(\"mask\", None)\n # StoryMem\n if story_mem_latents is not None:\n image_cond = torch.cat([story_mem_latents.to(image_cond), image_cond], dim=1)\n image_cond_mask = torch.cat([torch.ones_like(story_mem_latents)[:4], image_cond_mask], dim=1) if image_cond_mask is not None else None\n\n if image_cond_mask is not None:\n image_cond = torch.cat([image_cond_mask, image_cond])\n else:\n image_cond[:, 1:] = 0\n\n # ATI tracks (not bundled — use Bernini Context Embeds instead)\n if transformer_options is not None:\n ATI_tracks = transformer_options.get(\"ati_tracks\", None)\n if ATI_tracks is not None:\n raise ValueError(\n \"ATI motion tracks are not supported in ComfyUI-Bernini. \"\n \"Use Bernini Context Embeds for reference conditioning.\"\n )\n\n add_cond_latents = image_embeds.get(\"add_cond_latents\", None)\n if add_cond_latents is not None:\n add_cond = add_cond_latents[\"pose_latent\"]\n attn_cond = add_cond_latents[\"ref_latent\"]\n attn_cond_neg = add_cond_latents[\"ref_latent_neg\"]\n add_cond_start_percent = add_cond_latents[\"pose_cond_start_percent\"]\n add_cond_end_percent = add_cond_latents[\"pose_cond_end_percent\"]\n\n end_image = image_embeds.get(\"end_image\", None)\n fun_or_fl2v_model = image_embeds.get(\"fun_or_fl2v_model\", False)\n latent_frames = (image_embeds[\"num_frames\"] - 1) // 4\n latent_frames = latent_frames + (2 if end_image is not None and not fun_or_fl2v_model else 1)\n latent_frames = latent_frames + story_mem_latents.shape[1] if story_mem_latents is not None else latent_frames\n noise = torch.randn( #C, T, H, W\n 48 if is_5b else 16,\n latent_frames,\n image_embeds[\"lat_h\"],\n image_embeds[\"lat_w\"],\n dtype=torch.float32,\n generator=seed_g,\n device=torch.device(\"cpu\"))\n\n seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1])\n\n control_embeds = image_embeds.get(\"control_embeds\", None)\n if control_embeds is not None:\n if transformer.in_dim not in [148, 52, 48, 36, 32]:\n raise ValueError(\"Control signal only works with Fun-Control model\")\n\n control_latents = control_embeds.get(\"control_images\", None)\n control_start_percent = control_embeds.get(\"start_percent\", 0.0)\n control_end_percent = control_embeds.get(\"end_percent\", 1.0)\n control_camera_latents = control_embeds.get(\"control_camera_latents\", None)\n if control_camera_latents is not None:\n if transformer.control_adapter is None:\n raise ValueError(\"Control camera latents are only supported with Fun-Control-Camera model\")\n control_camera_start_percent = control_embeds.get(\"control_camera_start_percent\", 0.0)\n control_camera_end_percent = control_embeds.get(\"control_camera_end_percent\", 1.0)\n\n drop_last = image_embeds.get(\"drop_last\", False)\n else: #t2v\n target_shape = image_embeds.get(\"target_shape\", None)\n if target_shape is None:\n raise ValueError(\"Empty image embeds must be provided for T2V models\")\n\n # VACE\n vace_context = image_embeds.get(\"vace_context\", None)\n vace_scale = image_embeds.get(\"vace_scale\", None)\n if not isinstance(vace_scale, list):\n vace_scale = [vace_scale] * (steps+1)\n vace_start_percent = image_embeds.get(\"vace_start_percent\", 0.0)\n vace_end_percent = image_embeds.get(\"vace_end_percent\", 1.0)\n vace_seqlen = image_embeds.get(\"vace_seq_len\", None)\n\n vace_additional_embeds = image_embeds.get(\"additional_vace_inputs\", [])\n if vace_context is not None:\n vace_data = [\n {\"context\": vace_context,\n \"scale\": vace_scale,\n \"start\": vace_start_percent,\n \"end\": vace_end_percent,\n \"seq_len\": vace_seqlen\n }\n ]\n if len(vace_additional_embeds) > 0:\n for i in range(len(vace_additional_embeds)):\n if vace_additional_embeds[i].get(\"has_ref\", False):\n has_ref = True\n vace_scale = vace_additional_embeds[i][\"vace_scale\"]\n if not isinstance(vace_scale, list):\n vace_scale = [vace_scale] * (steps+1)\n vace_data.append({\n \"context\": vace_additional_embeds[i][\"vace_context\"],\n \"scale\": vace_scale,\n \"start\": vace_additional_embeds[i][\"vace_start_percent\"],\n \"end\": vace_additional_embeds[i][\"vace_end_percent\"],\n \"seq_len\": vace_additional_embeds[i][\"vace_seq_len\"]\n })\n\n noise = torch.randn(\n 48 if is_5b else 16,\n target_shape[1] + 1 if has_ref else target_shape[1],\n target_shape[2] // 2 if is_5b else target_shape[2], #todo make this smarter\n target_shape[3] // 2 if is_5b else target_shape[3], #todo make this smarter\n dtype=torch.float32,\n device=torch.device(\"cpu\"),\n generator=seed_g)\n\n seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1])\n\n recammaster = image_embeds.get(\"recammaster\", None)\n if recammaster is not None:\n camera_embed = recammaster.get(\"camera_embed\", None)\n recam_latents = recammaster.get(\"source_latents\", None)\n orig_noise_len = noise.shape[1]\n log.info(f\"RecamMaster camera embed shape: {camera_embed.shape}\")\n log.info(f\"RecamMaster source video shape: {recam_latents.shape}\")\n seq_len *= 2\n\n if image_embeds.get(\"mocha_embeds\", None) is not None:\n mocha_embeds = image_embeds.get(\"mocha_embeds\", None)\n mocha_num_refs = image_embeds.get(\"mocha_num_refs\", 0)\n orig_noise_len = noise.shape[1]\n seq_len = image_embeds.get(\"seq_len\", seq_len)\n log.info(f\"MoCha embeds shape: {mocha_embeds.shape}\")\n\n # Fun control and control lora\n control_embeds = image_embeds.get(\"control_embeds\", None)\n if control_embeds is not None:\n control_latents = control_embeds.get(\"control_images\", None)\n if control_latents is not None:\n control_latents = control_latents.to(device)\n\n control_camera_latents = control_embeds.get(\"control_camera_latents\", None)\n if control_camera_latents is not None:\n if transformer.control_adapter is None:\n raise ValueError(\"Control camera latents are only supported with Fun-Control-Camera model\")\n control_camera_start_percent = control_embeds.get(\"control_camera_start_percent\", 0.0)\n control_camera_end_percent = control_embeds.get(\"control_camera_end_percent\", 1.0)\n\n if control_lora:\n image_cond = control_latents.to(device)\n if not patcher.model.is_patched:\n log.info(\"Re-loading control LoRA...\")\n patcher = apply_lora(patcher, device, device, low_mem_load=False, control_lora=True)\n patcher.model.is_patched = True\n else:\n if transformer.in_dim not in [148, 48, 36, 32, 52]:\n raise ValueError(\"Control signal only works with Fun-Control model\")\n image_cond = torch.zeros_like(noise).to(device) #fun control\n if transformer.in_dim in [148, 52] or transformer.control_adapter is not None: #fun 2.2 control\n mask_latents = torch.tile(\n torch.zeros_like(noise[:1]), [4, 1, 1, 1]\n )\n masked_video_latents_input = torch.zeros_like(noise)\n image_cond = torch.cat([mask_latents, masked_video_latents_input], dim=0).to(device)\n clip_fea = None\n fun_ref_image = control_embeds.get(\"fun_ref_image\", None)\n if fun_ref_image is not None:\n if transformer.ref_conv.weight.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:\n raise ValueError(\"Fun-Control reference image won't work with this specific fp8_scaled model, it's been fixed in latest version of the model\")\n control_start_percent = control_embeds.get(\"start_percent\", 0.0)\n control_end_percent = control_embeds.get(\"end_percent\", 1.0)\n else:\n if transformer.in_dim in [148, 52]: #fun inp\n mask_latents = torch.tile(\n torch.zeros_like(noise[:1]), [4, 1, 1, 1]\n )\n masked_video_latents_input = torch.zeros_like(noise)\n image_cond = torch.cat([mask_latents, masked_video_latents_input], dim=0).to(device)\n\n # Phantom inputs\n phantom_latents = image_embeds.get(\"phantom_latents\", None)\n phantom_cfg_scale = image_embeds.get(\"phantom_cfg_scale\", None)\n if not isinstance(phantom_cfg_scale, list):\n phantom_cfg_scale = [phantom_cfg_scale] * (steps +1)\n phantom_start_percent = image_embeds.get(\"phantom_start_percent\", 0.0)\n phantom_end_percent = image_embeds.get(\"phantom_end_percent\", 1.0)\n\n # CLIP image features\n clip_fea = image_embeds.get(\"clip_context\", None)\n if clip_fea is not None:\n clip_fea = clip_fea.to(dtype)\n clip_fea_neg = image_embeds.get(\"negative_clip_context\", None)\n if clip_fea_neg is not None:\n clip_fea_neg = clip_fea_neg.to(dtype)\n\n num_frames = image_embeds.get(\"num_frames\", 0)\n\n # HuMo (removed — not used in Bernini workflows)\n humo_audio = humo_audio_neg = humo_image_cond = humo_image_cond_neg = None\n humo_reference_count = 0\n humo_audio_scale = 1.0\n humo_audio_cfg_scale = 1.0\n humo_start_percent = 0.0\n humo_end_percent = 1.0\n if not isinstance(humo_audio_cfg_scale, list):\n humo_audio_cfg_scale = [humo_audio_cfg_scale] * (steps + 1)\n\n pos_latent = neg_latent = None\n\n # Ovi\n noise_audio = latent_ovi = seq_len_ovi = None\n if transformer.audio_model is not None:\n noise_audio = samples.get(\"latent_ovi_audio\", None) if samples is not None else None\n if noise_audio is not None:\n if not torch.any(noise_audio):\n noise_audio = torch.randn(noise_audio.shape, device=torch.device(\"cpu\"), dtype=torch.float32, generator=seed_g)\n else:\n noise_audio = noise_audio.squeeze().movedim(0, 1).to(device, dtype)\n else:\n noise_audio = torch.randn((157, 20), device=torch.device(\"cpu\"), dtype=torch.float32, generator=seed_g) # T C\n log.info(f\"Ovi audio latent shape: {noise_audio.shape}\")\n latent_ovi = noise_audio\n seq_len_ovi = noise_audio.shape[0]\n\n from .bernini_unsupported import reject_pose_face_inputs\n reject_pose_face_inputs(image_embeds)\n\n context_latents = image_embeds.get(\"context_latents\", None)\n\n if image_cond is None:\n image_cond = image_embeds.get(\"ref_latent\", None)\n has_ref = image_cond is not None or has_ref\n\n latent_video_length = noise.shape[1]\n\n if freeinit_args is not None:\n raise ValueError(\n \"FreeInit is not supported in ComfyUI-Bernini. \"\n \"Remove freeinit_args from the sampler extra inputs.\"\n )\n\n # UniAnimate\n if unianimate_poses is not None:\n transformer.dwpose_embedding.to(device, dtype)\n dwpose_data = unianimate_poses[\"pose\"].to(device, dtype)\n dwpose_data = torch.cat([dwpose_data[:,:,:1].repeat(1,1,3,1,1), dwpose_data], dim=2)\n dwpose_data = transformer.dwpose_embedding(dwpose_data)\n log.info(f\"UniAnimate pose embed shape: {dwpose_data.shape}\")\n if not multitalk_sampling:\n if dwpose_data.shape[2] > latent_video_length:\n log.warning(f\"UniAnimate pose embed length {dwpose_data.shape[2]} is longer than the video length {latent_video_length}, truncating\")\n dwpose_data = dwpose_data[:,:, :latent_video_length]\n elif dwpose_data.shape[2] < latent_video_length:\n log.warning(f\"UniAnimate pose embed length {dwpose_data.shape[2]} is shorter than the video length {latent_video_length}, padding with last pose\")\n pad_len = latent_video_length - dwpose_data.shape[2]\n pad = dwpose_data[:,:,:1].repeat(1,1,pad_len,1,1)\n dwpose_data = torch.cat([dwpose_data, pad], dim=2)\n\n random_ref_dwpose_data = None\n if image_cond is not None:\n transformer.randomref_embedding_pose.to(device, dtype)\n random_ref_dwpose = unianimate_poses.get(\"ref\", None)\n if random_ref_dwpose is not None:\n random_ref_dwpose_data = transformer.randomref_embedding_pose(\n random_ref_dwpose.to(device, dtype)\n ).unsqueeze(2).to(dtype) # [1, 20, 104, 60]\n del random_ref_dwpose\n\n unianim_data = {\n \"dwpose\": dwpose_data,\n \"random_ref\": random_ref_dwpose_data.squeeze(0) if random_ref_dwpose_data is not None else None,\n \"strength\": unianimate_poses[\"strength\"],\n \"start_percent\": unianimate_poses[\"start_percent\"],\n \"end_percent\": unianimate_poses[\"end_percent\"]\n }\n\n # FantasyTalking\n audio_proj = multitalk_audio_embeds = None\n audio_scale = 1.0\n if fantasytalking_embeds is not None:\n audio_proj = fantasytalking_embeds[\"audio_proj\"].to(device)\n audio_scale = fantasytalking_embeds[\"audio_scale\"]\n audio_cfg_scale = fantasytalking_embeds[\"audio_cfg_scale\"]\n if not isinstance(audio_cfg_scale, list):\n audio_cfg_scale = [audio_cfg_scale] * (steps +1)\n log.info(f\"Audio proj shape: {audio_proj.shape}\")\n\n\n # MultiTalk\n multitalk_audio_embeds = audio_emb_slice = audio_features_in = None\n multitalk_audio_stride = None\n multitalk_embeds = image_embeds.get(\"multitalk_embeds\", multitalk_embeds)\n\n if multitalk_embeds is not None:\n audio_emb_slice = multitalk_embeds.get(\"audio_emb_slice\", None) # if already sliced\n # Handle single or multiple speaker embeddings\n if audio_emb_slice is None:\n audio_features_in = multitalk_embeds.get(\"audio_features\", None)\n if audio_features_in is not None:\n if isinstance(audio_features_in, list):\n multitalk_audio_embeds = [emb.to(device, dtype) for emb in audio_features_in]\n else:\n # keep backward-compatibility with single tensor input\n multitalk_audio_embeds = [audio_features_in.to(device, dtype)]\n\n shapes = [tuple(e.shape) for e in multitalk_audio_embeds]\n log.info(f\"Multitalk audio features shapes (per speaker): {shapes}\")\n\n audio_scale = multitalk_embeds.get(\"audio_scale\", 1.0)\n audio_cfg_scale = multitalk_embeds.get(\"audio_cfg_scale\", 1.0)\n ref_target_masks = multitalk_embeds.get(\"ref_target_masks\", None)\n if not isinstance(audio_cfg_scale, list):\n audio_cfg_scale = [audio_cfg_scale] * (steps + 1)\n\n # FantasyPortrait\n fantasy_portrait_input = None\n fantasy_portrait_embeds = image_embeds.get(\"portrait_embeds\", None)\n if fantasy_portrait_embeds is not None:\n log.info(\"Using FantasyPortrait embeddings\")\n fantasy_portrait_input = fantasy_portrait_embeds.copy()\n portrait_cfg = fantasy_portrait_input.get(\"cfg_scale\", 1.0)\n if not isinstance(portrait_cfg, list):\n portrait_cfg = [portrait_cfg] * (steps + 1)\n\n # MiniMax Remover\n minimax_latents = image_embeds.get(\"minimax_latents\", None)\n minimax_mask_latents = image_embeds.get(\"minimax_mask_latents\", None)\n if minimax_latents is not None:\n log.info(f\"minimax_latents: {minimax_latents.shape}, minimax_mask_latents: {minimax_mask_latents.shape}\")\n minimax_latents = minimax_latents.to(device, dtype)\n minimax_mask_latents = minimax_mask_latents.to(device, dtype)\n\n # Context windows\n is_looped = False\n context_reference_latent = None\n if context_options is not None:\n if context_options[\"context_frames\"] <= num_frames:\n context_schedule = context_options[\"context_schedule\"]\n context_frames = (context_options[\"context_frames\"] - 1) // 4 + 1\n context_stride = context_options[\"context_stride\"] // 4\n context_overlap = context_options[\"context_overlap\"] // 4\n context_reference_latent = context_options.get(\"reference_latent\", None)\n\n # Get total number of prompts\n num_prompts = len(text_embeds[\"prompt_embeds\"])\n log.info(f\"Number of prompts: {num_prompts}\")\n # Calculate which section this context window belongs to\n section_size = (latent_video_length / num_prompts) if num_prompts != 0 else 1\n log.info(f\"Section size: {section_size}\")\n is_looped = context_schedule == \"uniform_looped\"\n\n if mocha_embeds is not None:\n seq_len = (context_frames * 2 + 1 + mocha_num_refs) * (noise.shape[2] * noise.shape[3] // 4)\n else:\n seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * context_frames)\n log.info(f\"context window seq len: {seq_len}\")\n\n if context_options[\"freenoise\"]:\n log.info(\"Applying FreeNoise\")\n # code from AnimateDiff-Evolved by Kosinkadink (https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved)\n delta = context_frames - context_overlap\n for start_idx in range(0, latent_video_length-context_frames, delta):\n place_idx = start_idx + context_frames\n if place_idx >= latent_video_length:\n break\n end_idx = place_idx - 1\n\n if end_idx + delta >= latent_video_length:\n final_delta = latent_video_length - place_idx\n list_idx = torch.tensor(list(range(start_idx,start_idx+final_delta)), device=torch.device(\"cpu\"), dtype=torch.long)\n list_idx = list_idx[torch.randperm(final_delta, generator=seed_g)]\n noise[:, place_idx:place_idx + final_delta, :, :] = noise[:, list_idx, :, :]\n break\n list_idx = torch.tensor(list(range(start_idx,start_idx+delta)), device=torch.device(\"cpu\"), dtype=torch.long)\n list_idx = list_idx[torch.randperm(delta, generator=seed_g)]\n noise[:, place_idx:place_idx + delta, :, :] = noise[:, list_idx, :, :]\n\n log.info(f\"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap\")\n from .context_windows.context import get_context_scheduler, create_window_mask, WindowTracker\n self.window_tracker = WindowTracker(verbose=context_options[\"verbose\"])\n context = get_context_scheduler(context_schedule)\n else:\n log.info(\"Context frames is larger than total num_frames, disabling context windows\")\n context_options = None\n\n # MTV Crafter (removed — not used in Bernini workflows)\n mtv_input = None\n mtv_motion_tokens = None\n\n #region S2V\n s2v_audio_input = s2v_ref_latent = s2v_pose = s2v_ref_motion = None\n framepack = False\n s2v_audio_embeds = image_embeds.get(\"audio_embeds\", None)\n if s2v_audio_embeds is not None:\n log.info(\"Using S2V audio embeddings\")\n framepack = s2v_audio_embeds.get(\"enable_framepack\", False)\n if framepack and context_options is not None:\n raise ValueError(\"S2V framepack and context windows cannot be used at the same time\")\n\n s2v_audio_input = s2v_audio_embeds.get(\"audio_embed_bucket\", None)\n if s2v_audio_input is not None:\n #s2v_audio_input = s2v_audio_input[..., 0:image_embeds[\"num_frames\"]]\n s2v_audio_input = s2v_audio_input.to(device, dtype)\n s2v_audio_scale = s2v_audio_embeds[\"audio_scale\"]\n s2v_ref_latent = s2v_audio_embeds.get(\"ref_latent\", None)\n if s2v_ref_latent is not None:\n s2v_ref_latent = s2v_ref_latent.to(device, dtype)\n s2v_ref_motion = s2v_audio_embeds.get(\"ref_motion\", None)\n if s2v_ref_motion is not None:\n s2v_ref_motion = s2v_ref_motion.to(device, dtype)\n s2v_pose = s2v_audio_embeds.get(\"pose_latent\", None)\n if s2v_pose is not None:\n s2v_pose = s2v_pose.to(device, dtype)\n s2v_pose_start_percent = s2v_audio_embeds.get(\"pose_start_percent\", 0.0)\n s2v_pose_end_percent = s2v_audio_embeds.get(\"pose_end_percent\", 1.0)\n s2v_num_repeat = s2v_audio_embeds.get(\"num_repeat\", 1)\n vae = s2v_audio_embeds.get(\"vae\", None)\n\n # vid2vid\n noise_mask=original_image=None\n input_noise_mask = None\n if samples is not None and not multitalk_sampling:\n saved_generator_state = samples.get(\"generator_state\", None)\n if saved_generator_state is not None:\n seed_g.set_state(saved_generator_state)\n input_samples = samples.get(\"samples\", None)\n if input_samples is not None:\n input_samples = input_samples.squeeze(0).to(noise)\n if input_samples.shape[1] != noise.shape[1]:\n input_samples = torch.cat([input_samples[:, :1].repeat(1, noise.shape[1] - input_samples.shape[1], 1, 1), input_samples], dim=1)\n\n raw_noise_mask = samples.get(\"noise_mask\", None)\n if add_noise_to_samples:\n latent_timestep = timesteps[:1].to(noise)\n blended = noise * latent_timestep / 1000 + (1 - latent_timestep / 1000) * input_samples\n if raw_noise_mask is not None:\n # SCAIL continuity: only blend prefix from input_samples; suffix keeps full noise.\n nm = _latent_noise_mask_to_thw(raw_noise_mask.to(noise.device))\n if nm.shape[0] < noise.shape[1]:\n nm = nm.repeat(noise.shape[1] // nm.shape[0], 1, 1)\n nm = torch.nn.functional.interpolate(\n nm.unsqueeze(0).unsqueeze(0),\n size=(noise.shape[1], noise.shape[2], noise.shape[3]),\n mode='trilinear',\n align_corners=False,\n ).squeeze(0).squeeze(0)\n lock = (1.0 - nm).clamp(0.0, 1.0).unsqueeze(0).expand_as(noise)\n noise = blended * lock + noise * (1.0 - lock)\n else:\n noise = blended\n else:\n noise = input_samples\n\n noise_mask = raw_noise_mask\n if noise_mask is not None:\n input_noise_mask = noise_mask\n log.info(f\"Latent noise_mask shape: {noise_mask.shape}\")\n original_image = samples.get(\"original_image\", None)\n if original_image is None:\n original_image = input_samples\n noise_mask = _latent_noise_mask_to_thw(noise_mask)\n if noise_mask.shape[0] < noise.shape[1]:\n noise_mask = noise_mask.repeat(noise.shape[1] // noise_mask.shape[0], 1, 1)\n\n noise_mask = torch.nn.functional.interpolate(\n noise_mask.unsqueeze(0).unsqueeze(0), # Add batch and channel dims [1,1,T,H,W]\n size=(noise.shape[1], noise.shape[2], noise.shape[3]),\n mode='trilinear',\n align_corners=False\n ).repeat(1, noise.shape[0], 1, 1, 1)\n\n # extra latents (Pusa) and 5b\n latents_to_insert = add_index = noise_multipliers = None\n extra_latents = image_embeds.get(\"extra_latents\", None)\n clean_latent_indices = []\n noise_multiplier_list = image_embeds.get(\"pusa_noise_multipliers\", None)\n if noise_multiplier_list is not None:\n if len(noise_multiplier_list) != latent_video_length:\n noise_multipliers = torch.zeros(latent_video_length)\n else:\n noise_multipliers = torch.tensor(noise_multiplier_list)\n log.info(f\"Using Pusa noise multipliers: {noise_multipliers}\")\n if extra_latents is not None and transformer.multitalk_model_type.lower() != \"infinitetalk\":\n if noise_multiplier_list is not None:\n noise_multiplier_list = list(noise_multiplier_list) + [1.0] * (len(clean_latent_indices) - len(noise_multiplier_list))\n for i, entry in enumerate(extra_latents):\n add_index = entry[\"index\"]\n num_extra_frames = entry[\"samples\"].shape[2]\n # Handle negative indices\n if add_index < 0:\n add_index = noise.shape[1] + add_index\n add_index = max(0, min(add_index, noise.shape[1] - num_extra_frames))\n if start_step == 0:\n noise[:, add_index:add_index+num_extra_frames] = entry[\"samples\"].to(noise)\n log.info(f\"Adding extra samples to latent indices {add_index} to {add_index+num_extra_frames-1}\")\n clean_latent_indices.extend(range(add_index, add_index+num_extra_frames))\n if noise_multipliers is not None and len(noise_multiplier_list) != latent_video_length:\n for i, idx in enumerate(clean_latent_indices):\n noise_multipliers[idx] = noise_multiplier_list[i]\n log.info(f\"Using Pusa noise multipliers: {noise_multipliers}\")\n\n # lucy edit\n extra_channel_latents = image_embeds.get(\"extra_channel_latents\", None)\n if extra_channel_latents is not None:\n extra_channel_latents = extra_channel_latents[0].to(noise)\n\n # FlashVSR\n flashvsr_LQ_latent = LQ_images = None\n flashvsr_LQ_images = image_embeds.get(\"flashvsr_LQ_images\", None)\n flashvsr_strength = image_embeds.get(\"flashvsr_strength\", 1.0)\n if flashvsr_LQ_images is not None:\n flashvsr_LQ_images = flashvsr_LQ_images[:num_frames]\n first_frame = flashvsr_LQ_images[:1]\n last_frame = flashvsr_LQ_images[-1:].repeat(3, 1, 1, 1)\n flashvsr_LQ_images = torch.cat([first_frame, flashvsr_LQ_images, last_frame], dim=0)\n LQ_images = flashvsr_LQ_images.unsqueeze(0).movedim(-1, 1).to(dtype) * 2 - 1\n if context_options is None:\n flashvsr_LQ_latent = transformer.LQ_proj_in(LQ_images.to(device))\n log.info(f\"flashvsr_LQ_latent: {flashvsr_LQ_latent[0].shape}\")\n seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1])\n\n latent = noise\n\n # LongCat-Avatar\n longcat_ref_latent = None\n longcat_num_ref_latents = longcat_num_cond_latents = 0\n longcat_avatar_options = image_embeds.get(\"longcat_avatar_options\", None)\n\n if longcat_avatar_options is not None:\n longcat_ref_latent = longcat_avatar_options.get(\"longcat_ref_latent\", None)\n if longcat_ref_latent is not None:\n log.info(f\"LongCat-Avatar reference latent shape: {longcat_ref_latent.shape}\")\n latent = torch.cat([longcat_ref_latent.to(latent), latent], dim=1)\n seq_len = math.ceil((latent.shape[2] * latent.shape[3]) / 4 * latent.shape[1])\n insert_len = longcat_ref_latent.shape[1]\n clean_latent_indices = list(range(0, insert_len)) + [i + insert_len for i in clean_latent_indices]\n longcat_num_ref_latents = longcat_ref_latent.shape[1]\n latent_video_length += insert_len\n longcat_num_cond_latents = len(clean_latent_indices)\n log.info(f\"LongCat num_cond_latents: {longcat_num_cond_latents} num_ref_latents: {longcat_num_ref_latents}\")\n # v1.5 (Whisper) embeds set audio_stride=1; v1.0 (wav2vec2) uses 2 for LongCat\n if multitalk_audio_stride is not None:\n audio_stride = multitalk_audio_stride\n else:\n audio_stride = 2 if transformer.is_longcat else 1\n\n #controlnet\n controlnet_latents = controlnet = None\n if transformer_options is not None:\n controlnet = transformer_options.get(\"controlnet\", None)\n if controlnet is not None:\n self.controlnet = controlnet[\"controlnet\"]\n controlnet_start = controlnet[\"controlnet_start\"]\n controlnet_end = controlnet[\"controlnet_end\"]\n controlnet_latents = controlnet[\"control_latents\"]\n controlnet[\"controlnet_weight\"] = controlnet[\"controlnet_strength\"]\n controlnet[\"controlnet_stride\"] = controlnet[\"control_stride\"]\n\n #uni3c\n uni3c_data = uni3c_data_input = None\n if uni3c_embeds is not None:\n transformer.uni3c_controlnet = uni3c_embeds[\"controlnet\"]\n render_latent = uni3c_embeds[\"render_latent\"].to(device)\n uni3c_data = uni3c_embeds.copy()\n if render_latent.shape != noise.shape:\n render_latent = torch.nn.functional.interpolate(\n render_latent, size=(noise.shape[1], noise.shape[2], noise.shape[3]),\n mode='trilinear', align_corners=False,\n )\n uni3c_data[\"render_latent\"] = render_latent\n\n # Enhance-a-video (feta)\n if feta_args is not None and latent_video_length > 1:\n set_enhance_weight(feta_args[\"weight\"])\n feta_start_percent = feta_args[\"start_percent\"]\n feta_end_percent = feta_args[\"end_percent\"]\n set_num_frames(latent_video_length) if context_options is None else set_num_frames(context_frames)\n enhance_enabled = True\n else:\n feta_args = None\n enhance_enabled = False\n\n # Bindweave\n qwenvl_embeds_pos = image_embeds.get(\"qwenvl_embeds_pos\", None)\n qwenvl_embeds_neg = image_embeds.get(\"qwenvl_embeds_neg\", None)\n\n mm.unload_all_models()\n mm.soft_empty_cache()\n gc.collect()\n\n #blockswap init\n init_blockswap(transformer, block_swap_args, model)\n\n # Initialize Cache if enabled\n previous_cache_states = None\n transformer.enable_teacache = transformer.enable_magcache = transformer.enable_easycache = False\n cache_args = teacache_args if teacache_args is not None else cache_args #for backward compatibility on old workflows\n if cache_args is None and teacache_args is None and image_embeds.get(\"bernini_pipeline\"):\n if enable_teacache is True:\n from .bernini_perf import default_teacache_args\n\n cache_args = default_teacache_args(timesteps)\n log.info(\"Bernini perf: TeaCache enabled (rel_l1_thresh=0.15, mode=e, coefficients on)\")\n if cache_args is not None:\n from .cache_methods.cache_methods import set_transformer_cache_method\n transformer = set_transformer_cache_method(transformer, timesteps, cache_args)\n\n # Initialize cache state\n if samples is not None:\n previous_cache_states = samples.get(\"cache_states\", None)\n if previous_cache_states is not None:\n log.info(\"Using cache states from previous sampler\")\n self.cache_state = previous_cache_states[\"cache_state\"]\n transformer.easycache_state = previous_cache_states[\"easycache_state\"]\n transformer.magcache_state = previous_cache_states[\"magcache_state\"]\n transformer.teacache_state = previous_cache_states[\"teacache_state\"]\n\n if previous_cache_states is None:\n self.cache_state = [None, None]\n if phantom_latents is not None:\n log.info(f\"Phantom latents shape: {phantom_latents.shape}\")\n self.cache_state = [None, None, None]\n self.cache_state_source = [None, None]\n self.cache_states_context = []\n\n # Skip layer guidance (SLG)\n if slg_args is not None:\n assert batched_cfg is not None, \"Batched cfg is not supported with SLG\"\n transformer.slg_blocks = slg_args[\"blocks\"]\n transformer.slg_start_percent = slg_args[\"start_percent\"]\n transformer.slg_end_percent = slg_args[\"end_percent\"]\n else:\n transformer.slg_blocks = None\n\n # Setup radial attention\n if transformer.attention_mode == \"radial_sage_attention\":\n setup_radial_attention(transformer, transformer_options, latent, seq_len, latent_video_length, context_options=context_options)\n\n # Experimental args\n use_cfg_zero_star = use_tangential = use_fresca = bidirectional_sampling = use_tsr = False\n raag_alpha = 0.0\n transformer.video_attention_split_steps = []\n if experimental_args is not None:\n video_attention_split_steps = experimental_args.get(\"video_attention_split_steps\", [])\n if video_attention_split_steps:\n transformer.video_attention_split_steps = [int(x.strip()) for x in video_attention_split_steps.split(\",\")]\n\n use_zero_init = experimental_args.get(\"use_zero_init\", True)\n use_cfg_zero_star = experimental_args.get(\"cfg_zero_star\", False)\n use_tangential = experimental_args.get(\"use_tcfg\", False)\n zero_star_steps = experimental_args.get(\"zero_star_steps\", 0)\n raag_alpha = experimental_args.get(\"raag_alpha\", 0.0)\n\n use_fresca = experimental_args.get(\"use_fresca\", False)\n if use_fresca:\n fresca_scale_low = experimental_args.get(\"fresca_scale_low\", 1.0)\n fresca_scale_high = experimental_args.get(\"fresca_scale_high\", 1.25)\n fresca_freq_cutoff = experimental_args.get(\"fresca_freq_cutoff\", 20)\n\n bidirectional_sampling = experimental_args.get(\"bidirectional_sampling\", False)\n if bidirectional_sampling:\n sample_scheduler_flipped = copy.deepcopy(sample_scheduler)\n use_tsr = experimental_args.get(\"temporal_score_rescaling\", False)\n tsr_k = experimental_args.get(\"tsr_k\", 1.0)\n tsr_sigma = experimental_args.get(\"tsr_sigma\", 1.0)\n\n # Rotary positional embeddings (RoPE)\n\n # RoPE base freq scaling as used with CineScale\n ntk_alphas = [1.0, 1.0, 1.0]\n if isinstance(rope_function, dict):\n ntk_alphas = rope_function[\"ntk_scale_f\"], rope_function[\"ntk_scale_h\"], rope_function[\"ntk_scale_w\"]\n rope_function = rope_function[\"rope_function\"]\n\n # Stand-In\n standin_input = image_embeds.get(\"standin_input\", None)\n if standin_input is not None:\n rope_function = \"comfy\" # only works with this currently\n\n freqs = None\n\n log.info(f\"Rope function: {rope_function}\")\n\n riflex_freq_index = 0 if riflex_freq_index is None else riflex_freq_index\n transformer.rope_embedder.k = None\n transformer.rope_embedder.num_frames = None\n d = transformer.dim // transformer.num_heads\n\n if mocha_embeds is not None:\n raise ValueError(\n \"MoCha embeds are not supported in ComfyUI-Bernini. \"\n \"Use Bernini Context Embeds for reference conditioning.\"\n )\n elif \"default\" in rope_function or bidirectional_sampling: # original RoPE\n freqs = torch.cat([\n rope_params(1024, d - 4 * (d // 6), L_test=latent_video_length, k=riflex_freq_index),\n rope_params(1024, 2 * (d // 6)),\n rope_params(1024, 2 * (d // 6))\n ],\n dim=1)\n elif \"comfy\" in rope_function: # comfy's rope\n transformer.rope_embedder.k = riflex_freq_index\n transformer.rope_embedder.num_frames = latent_video_length\n\n transformer.rope_func = rope_function\n for block in transformer.blocks:\n block.rope_func = rope_function\n if transformer.vace_layers is not None:\n for block in transformer.vace_blocks:\n block.rope_func = rope_function\n\n # Lynx\n lynx_ref_buffer = None\n lynx_embeds = image_embeds.get(\"lynx_embeds\", None)\n if lynx_embeds is not None:\n if lynx_embeds.get(\"ip_x\", None) is not None:\n if transformer.blocks[0].cross_attn.ip_adapter is None:\n raise ValueError(\"Lynx IP embeds provided, but the no lynx ip adapter layers found in the model.\")\n lynx_embeds = lynx_embeds.copy()\n log.info(\"Using Lynx embeddings\", lynx_embeds)\n lynx_ref_latent = lynx_embeds.get(\"ref_latent\", None)\n lynx_ref_latent_uncond = lynx_embeds.get(\"ref_latent_uncond\", None)\n lynx_ref_text_embed = lynx_embeds.get(\"ref_text_embed\", None)\n lynx_ref_text_embed = dict_to_device(lynx_ref_text_embed, device)\n lynx_cfg_scale = lynx_embeds.get(\"cfg_scale\", 1.0)\n if not isinstance(lynx_cfg_scale, list):\n lynx_cfg_scale = [lynx_cfg_scale] * (steps + 1)\n\n if lynx_ref_latent is not None:\n if transformer.blocks[0].self_attn.ref_adapter is None:\n raise ValueError(\"Lynx reference provided, but the no lynx reference adapter layers found in the model.\")\n lynx_ref_latent = lynx_ref_latent[0]\n lynx_ref_latent_uncond = lynx_ref_latent_uncond[0]\n lynx_embeds[\"ref_feature_extractor\"] = True\n log.info(f\"Lynx ref latent shape: {lynx_ref_latent.shape}\")\n log.info(\"Extracting Lynx ref cond buffer...\")\n if transformer.in_dim == 36:\n mask_latents = torch.tile(torch.zeros_like(lynx_ref_latent[:1]), [4, 1, 1, 1])\n empty_image_cond = torch.cat([mask_latents, torch.zeros_like(lynx_ref_latent)], dim=0).to(device)\n lynx_ref_input = torch.cat([lynx_ref_latent, empty_image_cond], dim=0)\n else:\n lynx_ref_input = lynx_ref_latent\n lynx_ref_buffer = transformer(\n [lynx_ref_input.to(device, dtype)],\n torch.tensor([0], device=device),\n lynx_ref_text_embed[\"prompt_embeds\"],\n seq_len=math.ceil((lynx_ref_latent.shape[2] * lynx_ref_latent.shape[3]) / 4 * lynx_ref_latent.shape[1]),\n lynx_embeds=lynx_embeds\n )\n log.info(f\"Extracted {len(lynx_ref_buffer)} cond ref buffers\")\n if any(not math.isclose(c, 1.0) for c in cfg):\n log.info(\"Extracting Lynx ref uncond buffer...\")\n if transformer.in_dim == 36:\n lynx_ref_input_uncond = torch.cat([lynx_ref_latent_uncond, empty_image_cond], dim=0)\n else:\n lynx_ref_input_uncond = lynx_ref_latent_uncond\n lynx_ref_buffer_uncond = transformer(\n [lynx_ref_input_uncond.to(device, dtype)],\n torch.tensor([0], device=device),\n lynx_ref_text_embed[\"prompt_embeds\"],\n seq_len=math.ceil((lynx_ref_latent.shape[2] * lynx_ref_latent.shape[3]) / 4 * lynx_ref_latent.shape[1]),\n lynx_embeds=lynx_embeds,\n is_uncond=True\n )\n log.info(f\"Extracted {len(lynx_ref_buffer_uncond)} uncond ref buffers\")\n\n if lynx_embeds.get(\"ip_x\", None) is not None:\n lynx_embeds[\"ip_x\"] = lynx_embeds[\"ip_x\"].to(device, dtype)\n lynx_embeds[\"ip_x_uncond\"] = lynx_embeds[\"ip_x_uncond\"].to(device, dtype)\n lynx_embeds[\"ref_feature_extractor\"] = False\n lynx_embeds[\"ref_latent\"] = lynx_embeds[\"ref_text_embed\"] = None\n lynx_embeds[\"ref_buffer\"] = lynx_ref_buffer\n lynx_embeds[\"ref_buffer_uncond\"] = lynx_ref_buffer_uncond if not math.isclose(cfg[0], 1.0) else None\n mm.soft_empty_cache()\n\n # UniLumos\n foreground_latents = image_embeds.get(\"foreground_latents\", None)\n if foreground_latents is not None:\n log.info(f\"UniLumos foreground latent input shape: {foreground_latents.shape}\")\n foreground_latents = foreground_latents.to(device, dtype)\n background_latents = image_embeds.get(\"background_latents\", None)\n if background_latents is not None:\n log.info(f\"UniLumos background latent input shape: {background_latents.shape}\")\n background_latents = background_latents.to(device, dtype)\n\n #Time-to-move (TTM)\n ttm_start_step = 0\n ttm_reference_latents = image_embeds.get(\"ttm_reference_latents\", None)\n if ttm_reference_latents is not None:\n motion_mask = image_embeds[\"ttm_mask\"].to(device, dtype)\n ttm_start_step = max(image_embeds[\"ttm_start_step\"] - start_step, 0)\n ttm_end_step = image_embeds[\"ttm_end_step\"] - start_step\n\n if ttm_start_step > steps:\n raise ValueError(\"TTM start step is beyond the total number of steps\")\n\n if ttm_end_step > ttm_start_step:\n log.info(\"Using Time-to-move (TTM)\")\n log.info(f\"TTM reference latents shape: {ttm_reference_latents.shape}\")\n log.info(f\"TTM motion mask shape: {motion_mask.shape}\")\n log.info(f\"Applying TTM from step {ttm_start_step} to {ttm_end_step}\")\n\n latent = add_noise(ttm_reference_latents, noise, timesteps[ttm_start_step].to(noise.device)).to(latent)\n\n # SteadyDancer\n sdancer_embeds = image_embeds.get(\"sdancer_embeds\", None)\n sdancer_data = sdancer_input = None\n if sdancer_embeds is not None:\n log.info(\"Using SteadyDancer embeddings:\")\n for k, v in sdancer_embeds.items():\n log.info(f\" {k}: {v.shape if isinstance(v, torch.Tensor) else v}\")\n sdancer_data = sdancer_embeds.copy()\n sdancer_data = dict_to_device(sdancer_data, device, dtype)\n\n # One-to-all-Animation\n one_to_all_embeds = image_embeds.get(\"one_to_all_embeds\", None)\n one_to_all_data = prev_latents = None\n latents_to_not_step = 0\n if one_to_all_embeds is not None:\n log.info(\"Using One-to-All embeddings:\")\n for k, v in one_to_all_embeds.items():\n log.info(f\" {k}: {v.shape if isinstance(v, torch.Tensor) else v}\")\n one_to_all_data = one_to_all_embeds.copy()\n one_to_all_data = dict_to_device(one_to_all_data, device, dtype)\n if one_to_all_embeds.get(\"pose_images\") is not None:\n transformer.input_hint_block.to(device)\n pose_images_in = one_to_all_data.pop(\"pose_images\")\n pose_images = transformer.input_hint_block(pose_images_in)\n if one_to_all_embeds.get(\"ref_latent_pos\") is not None:\n pose_prefix_image = transformer.input_hint_block(one_to_all_data.pop(\"pose_prefix_image\"))\n pose_images = torch.cat([pose_prefix_image, pose_images],dim=2)\n one_to_all_data[\"controlnet_tokens\"] = pose_images.flatten(2).transpose(1, 2)\n transformer.input_hint_block.to(offload_device)\n\n one_to_all_pose_cfg_scale = one_to_all_embeds.get(\"pose_cfg_scale\", 1.0)\n if not isinstance(one_to_all_pose_cfg_scale, list):\n one_to_all_pose_cfg_scale = [one_to_all_pose_cfg_scale] * (steps + 1)\n\n prev_latents = one_to_all_data.get(\"prev_latents\", None)\n if prev_latents is not None:\n log.info(f\"Using previous latents for One-to-All Animation with shape: {prev_latents.shape}\")\n latent[:, :prev_latents.shape[1]] = prev_latents.to(latent)\n one_to_all_data[\"token_replace\"] = True\n latents_to_not_step = prev_latents.shape[1]\n one_to_all_data[\"num_latent_frames_to_replace\"] = latents_to_not_step\n\n # SCAIL\n scail_embeds = image_embeds.get(\"scail_embeds\", None)\n scail_data = None\n if scail_embeds is not None:\n log.info(\"Using SCAIL embeddings:\")\n for k, v in scail_embeds.items():\n log.info(f\" {k}: {v.shape if isinstance(v, torch.Tensor) else v}\")\n scail_data = scail_embeds.copy()\n scail_data = dict_to_device(scail_data, device, dtype)\n\n\n # WanMove (removed — not used in Bernini workflows)\n wanmove_embeds = None\n\n # LongVie2 dual control\n dual_control_embeds = image_embeds.get(\"dual_control\", None)\n if dual_control_embeds is not None and context_options is None:\n dual_control_input = dict_to_device(dual_control_embeds.copy(), device, dtype) if dual_control_embeds is not None else None\n prev_latents = dual_control_input.get(\"prev_latent\", None)\n if prev_latents is not None:\n _sigma = dual_control_embeds.get(\"first_frame_noise_level\", 0.925926)\n log.info(f\"Using dual control previous latents with first frame noise level: {_sigma}\")\n latent[:, :1] = (1 - _sigma) * prev_latents[:, -1:].to(latent) + _sigma * noise[:, :1]\n prev_ones = torch.ones(20, *prev_latents.shape[1:], device=device, dtype=dtype)\n dual_control_input[\"prev_latent\"] = torch.cat([prev_ones, prev_latents]).unsqueeze(0)\n\n # APG momentum buffer (persists across sampling steps)\n _momentum_buf = MomentumBuffer(apg_momentum) if uses_apg_guidance(guidance_mode) else None\n\n #region model pred\n def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, idx, image_cond=None, clip_fea=None,\n control_latents=None, vace_data=None, unianim_data=None, audio_proj=None, control_camera_latents=None,\n add_cond=None, cache_state=None, context_window=None, multitalk_audio_embeds=None, fantasy_portrait_input=None, reverse_time=False,\n mtv_motion_tokens=None, s2v_audio_input=None, s2v_ref_motion=None, s2v_motion_frames=[1, 0], s2v_pose=None,\n humo_image_cond=None, humo_image_cond_neg=None, humo_audio=None, humo_audio_neg=None,\n uni3c_data=None, latent_model_input_ovi=None, flashvsr_LQ_latent=None,\n context_latents=None, context_window_start=0,):\n nonlocal transformer\n nonlocal audio_cfg_scale\n\n autocast_enabled = (\"fp8\" in model[\"quantization\"] and not transformer.patched_linear)\n with torch.autocast(device_type=mm.get_autocast_device(device), dtype=dtype) if autocast_enabled else nullcontext():\n\n if use_cfg_zero_star and (idx <= zero_star_steps) and use_zero_init:\n return z*0, None\n\n nonlocal patcher\n current_step_percentage = idx / len(timesteps)\n control_lora_enabled = False\n image_cond_input = None\n if control_embeds is not None and control_camera_latents is None:\n if control_lora:\n control_lora_enabled = True\n else:\n if ((control_start_percent <= current_step_percentage <= control_end_percent) or \\\n (control_end_percent > 0 and idx == 0 and current_step_percentage >= control_start_percent)) and \\\n (control_latents is not None):\n image_cond_input = torch.cat([control_latents.to(z), image_cond.to(z)])\n else:\n image_cond_input = torch.cat([torch.zeros_like(noise, device=device, dtype=dtype), image_cond.to(z)])\n if fun_ref_image is not None:\n fun_ref_input = fun_ref_image.to(z)\n else:\n fun_ref_input = torch.zeros_like(z, dtype=z.dtype)[:, 0].unsqueeze(1)\n\n if control_lora:\n if not control_start_percent <= current_step_percentage <= control_end_percent:\n control_lora_enabled = False\n if patcher.model.is_patched:\n log.info(\"Unloading LoRA...\")\n patcher.unpatch_model(device)\n patcher.model.is_patched = False\n else:\n image_cond_input = control_latents.to(z)\n if not patcher.model.is_patched:\n log.info(\"Loading LoRA...\")\n patcher = apply_lora(patcher, device, device, low_mem_load=False, control_lora=True)\n patcher.model.is_patched = True\n\n elif ATI_tracks is not None and ((ati_start_percent <= current_step_percentage <= ati_end_percent) or\n (ati_end_percent > 0 and idx == 0 and current_step_percentage >= ati_start_percent)):\n image_cond_input = image_cond_ati.to(z)\n elif humo_image_cond is not None:\n humo_image_cond_neg_input = None\n if context_window is not None:\n image_cond_input = humo_image_cond[:, context_window].to(z)\n humo_image_cond_neg_input = humo_image_cond_neg[:, context_window].to(z)\n if humo_reference_count > 0:\n image_cond_input[:, -humo_reference_count:] = humo_image_cond[:, -humo_reference_count:]\n humo_image_cond_neg_input[:, -humo_reference_count:] = humo_image_cond_neg[:, -humo_reference_count:]\n else:\n if image_cond is not None:\n image_cond_input = image_cond.to(z)\n if humo_reference_count > 0:\n image_cond_input = torch.cat([image_cond_input, humo_image_cond[:, -humo_reference_count:].to(z)], dim=1)\n humo_image_cond_neg_input = torch.cat([image_cond_input, humo_image_cond_neg[:, -humo_reference_count:].to(z)], dim=1)\n else:\n image_cond_input = humo_image_cond.to(z)\n humo_image_cond_neg_input = humo_image_cond_neg.to(z)\n\n elif image_cond is not None:\n if reverse_time: # Flip the image condition\n image_cond_input = torch.cat([\n torch.flip(image_cond[:4], dims=[1]),\n torch.flip(image_cond[4:], dims=[1])\n ]).to(z)\n else:\n image_cond_input = image_cond.to(z)\n\n if control_camera_latents is not None:\n if (control_camera_start_percent <= current_step_percentage <= control_camera_end_percent) or \\\n (control_end_percent > 0 and idx == 0 and current_step_percentage >= control_camera_start_percent):\n control_camera_input = control_camera_latents.to(device, dtype)\n else:\n control_camera_input = None\n\n if recammaster is not None:\n z = torch.cat([z, recam_latents.to(z)], dim=1)\n\n if mocha_embeds is not None:\n if context_window is not None and mocha_embeds.shape[2] != context_frames:\n latent_frames = len(context_window)\n # [latent_frames, 1 mask frame, mocha_num_refs]\n latent_end = latent_frames\n mask_end = latent_end + 1\n partial_latents = mocha_embeds[:, context_window] # windowed latents\n mask_frame = mocha_embeds[:, latent_end:mask_end] # single mask frame\n ref_frames = mocha_embeds[:, -mocha_num_refs:] # reference frames\n\n partial_mocha_embeds = torch.cat([partial_latents, mask_frame, ref_frames], dim=1)\n z = torch.cat([z, partial_mocha_embeds.to(z)], dim=1)\n else:\n z = torch.cat([z, mocha_embeds.to(z)], dim=1)\n\n if mtv_input is not None:\n if ((mtv_start_percent <= current_step_percentage <= mtv_end_percent) or \\\n (mtv_end_percent > 0 and idx == 0 and current_step_percentage >= mtv_start_percent)):\n mtv_motion_tokens = mtv_motion_tokens.to(z)\n mtv_motion_rotary_emb = motion_rotary_emb\n\n use_phantom = False\n phantom_ref = None\n if phantom_latents is not None:\n if (phantom_start_percent <= current_step_percentage <= phantom_end_percent) or \\\n (phantom_end_percent > 0 and idx == 0 and current_step_percentage >= phantom_start_percent):\n phantom_ref = phantom_latents.to(z)\n use_phantom = True\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n\n if controlnet_latents is not None:\n if (controlnet_start <= current_step_percentage < controlnet_end):\n self.controlnet.to(device)\n controlnet_states = self.controlnet(\n hidden_states=z.unsqueeze(0).to(device, self.controlnet.dtype),\n timestep=timestep,\n encoder_hidden_states=positive_embeds[0].unsqueeze(0).to(device, self.controlnet.dtype),\n attention_kwargs=None,\n controlnet_states=controlnet_latents.to(device, self.controlnet.dtype),\n return_dict=False,\n )[0]\n if isinstance(controlnet_states, (tuple, list)):\n controlnet[\"controlnet_states\"] = [x.to(z) for x in controlnet_states]\n else:\n controlnet[\"controlnet_states\"] = controlnet_states.to(z)\n\n add_cond_input = None\n if add_cond is not None:\n if (add_cond_start_percent <= current_step_percentage <= add_cond_end_percent) or \\\n (add_cond_end_percent > 0 and idx == 0 and current_step_percentage >= add_cond_start_percent):\n add_cond_input = add_cond\n\n if minimax_latents is not None:\n if context_window is not None:\n z = torch.cat([z, minimax_latents[:, context_window], minimax_mask_latents[:, context_window]], dim=0)\n else:\n z = torch.cat([z, minimax_latents, minimax_mask_latents], dim=0)\n\n multitalk_audio_input = None\n if audio_emb_slice is not None:\n multitalk_audio_input = audio_emb_slice.to(z)\n elif not multitalk_sampling and multitalk_audio_embeds is not None:\n audio_embedding = multitalk_audio_embeds\n audio_embs = []\n indices = (torch.arange(4 + 1) - 2) * 1\n human_num = len(audio_embedding)\n # split audio with window size\n audio_end_idx = latent_video_length * 4 + 1 if add_cond is not None else (latent_video_length-1) * 4 + 1\n audio_end_idx = audio_end_idx * audio_stride\n if context_window is None:\n for human_idx in range(human_num):\n center_indices = torch.arange(0, audio_end_idx, audio_stride).unsqueeze(1) + indices.unsqueeze(0)\n center_indices = torch.clamp(center_indices, min=0, max=audio_embedding[human_idx].shape[0] - 1)\n\n audio_emb = audio_embedding[human_idx][center_indices].unsqueeze(0).to(device)\n audio_embs.append(audio_emb)\n else:\n for human_idx in range(human_num):\n audio_start = (context_window[0] * 4) * audio_stride\n audio_end = (context_window[-1] * 4 + 1) * audio_stride\n #print(\"audio_start: \", audio_start, \"audio_end: \", audio_end)\n center_indices = torch.arange(audio_start, audio_end, audio_stride).unsqueeze(1) + indices.unsqueeze(0)\n center_indices = torch.clamp(center_indices, min=0, max=audio_embedding[human_idx].shape[0] - 1)\n audio_emb = audio_embedding[human_idx][center_indices].unsqueeze(0).to(device)\n audio_embs.append(audio_emb)\n multitalk_audio_input = torch.concat(audio_embs, dim=0).to(dtype)\n\n elif multitalk_sampling and multitalk_audio_embeds is not None:\n multitalk_audio_input = multitalk_audio_embeds\n\n if context_window is not None and uni3c_data is not None and uni3c_data[\"render_latent\"].shape[2] != context_frames:\n uni3c_data_input = {\"render_latent\": uni3c_data[\"render_latent\"][:, :, context_window]}\n for k in uni3c_data:\n if k != \"render_latent\":\n uni3c_data_input[k] = uni3c_data[k]\n else:\n uni3c_data_input = uni3c_data\n\n if context_window is not None and sdancer_data is not None and sdancer_data[\"cond_pos\"].shape[1] != context_frames:\n sdancer_input = sdancer_data.copy()\n sdancer_input[\"cond_pos\"] = sdancer_data[\"cond_pos\"][:, context_window]\n sdancer_input[\"cond_neg\"] = sdancer_data[\"cond_neg\"][:, context_window] if sdancer_data.get(\"cond_neg\", None) is not None else None\n else:\n sdancer_input = sdancer_data\n\n if s2v_pose is not None:\n if not ((s2v_pose_start_percent <= current_step_percentage <= s2v_pose_end_percent) or \\\n (s2v_pose_end_percent > 0 and idx == 0 and current_step_percentage >= s2v_pose_start_percent)):\n s2v_pose = None\n\n\n if humo_audio is not None and ((humo_start_percent <= current_step_percentage <= humo_end_percent) or \\\n (humo_end_percent > 0 and idx == 0 and current_step_percentage >= humo_start_percent)):\n if context_window is None:\n humo_audio_input = humo_audio\n humo_audio_input_neg = humo_audio_neg if humo_audio_neg is not None else None\n else:\n humo_audio_input = humo_audio[context_window].to(z)\n if humo_audio_neg is not None:\n humo_audio_input_neg = humo_audio_neg[context_window].to(z)\n else:\n humo_audio_input_neg = None\n else:\n humo_audio_input = humo_audio_input_neg = None\n\n if extra_channel_latents is not None:\n if context_window is not None:\n extra_channel_latents_input = extra_channel_latents[:, context_window].to(z)\n else:\n extra_channel_latents_input = extra_channel_latents.to(z)\n z = torch.cat([z, extra_channel_latents_input])\n\n if \"rcm\" in sample_scheduler.__class__.__name__.lower():\n c_in = 1 / (torch.cos(timestep) + torch.sin(timestep))\n c_noise = (torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep))) * 1000\n z = z * c_in\n timestep = c_noise\n\n if image_cond is not None:\n self.noise_front_pad_num = image_cond_input.shape[1] - z.shape[1]\n if self.noise_front_pad_num > 0:\n pad = torch.zeros((z.shape[0], self.noise_front_pad_num, z.shape[2], z.shape[3]), dtype=z.dtype, device=z.device)\n z = torch.cat([pad, z], dim=1)\n nonlocal seq_len\n seq_len = math.ceil((z.shape[2] * z.shape[3]) / 4 * z.shape[1])\n\n if background_latents is not None or foreground_latents is not None:\n z = torch.cat([z, foreground_latents.to(z), background_latents.to(z)], dim=0)\n\n scail_data_in = None\n if scail_data is not None:\n ref_concat_mask = torch.zeros_like(z[:4])\n z = torch.cat([z, ref_concat_mask])\n if context_window is not None:\n scail_data_in = scail_data.copy()\n scail_data_in[\"pose_latent\"] = scail_data[\"pose_latent\"][:, context_window]\n else:\n scail_data_in = scail_data\n\n if wanmove_embeds is not None and context_window is not None:\n image_cond_input = replace_feature(image_cond_input.unsqueeze(0), track_pos[:, context_window].unsqueeze(0), wanmove_embeds.get(\"strength\", 1.0))[0]\n\n dual_control_in = None\n if dual_control_embeds is not None:\n if context_window is not None:\n dual_control_in = dual_control_embeds.copy()\n dense_input_latent = dual_control_embeds.get(\"dense_input_latent\", None)\n if dense_input_latent is not None:\n dual_control_in[\"dense_input_latent\"] = dual_control_embeds[\"dense_input_latent\"][:, :, context_window]\n sparse_input_latent = dual_control_embeds.get(\"sparse_input_latent\", None)\n if sparse_input_latent is not None:\n dual_control_in[\"sparse_input_latent\"] = dual_control_embeds[\"sparse_input_latent\"][:, :, context_window]\n else:\n dual_control_in = dual_control_input\n\n base_params = {\n 'x': [z], # latent\n 'y': [image_cond_input] if image_cond_input is not None else None, # image cond\n 'clip_fea': clip_fea, # clip features\n 'seq_len': seq_len, # sequence length\n 'device': device, # main device\n 'freqs': freqs, # rope freqs\n 't': timestep, # current timestep\n 'is_uncond': False, # is unconditional\n 'current_step': idx, # current step\n 'current_step_percentage': current_step_percentage, # current step percentage\n 'last_step': len(timesteps) - 1 == idx, # is last step\n 'control_lora_enabled': control_lora_enabled, # control lora toggle for patch embed selection\n 'enhance_enabled': enhance_enabled, # enhance-a-video toggle\n 'camera_embed': camera_embed, # recammaster embedding\n 'unianim_data': unianim_data, # unianimate input\n 'fun_ref': fun_ref_input if fun_ref_image is not None else None, # Fun model reference latent\n 'fun_camera': control_camera_input if control_camera_latents is not None else None, # Fun model camera embed\n 'audio_proj': audio_proj if fantasytalking_embeds is not None else None, # FantasyTalking audio projection\n 'audio_scale': audio_scale, # FantasyTalking audio scale\n \"uni3c_data\": uni3c_data_input, # Uni3C input\n \"controlnet\": controlnet, # TheDenk's controlnet input\n \"add_cond\": add_cond_input, # additional conditioning input\n \"nag_params\": text_embeds.get(\"nag_params\", {}), # normalized attention guidance\n \"nag_context\": text_embeds.get(\"nag_prompt_embeds\", None), # normalized attention guidance context\n \"multitalk_audio\": multitalk_audio_input, # Multi/InfiniteTalk audio input\n \"ref_target_masks\": ref_target_masks if multitalk_audio_embeds is not None else None, # Multi/InfiniteTalk reference target masks\n \"inner_t\": None,\n \"standin_input\": standin_input, # Stand-in reference input\n \"fantasy_portrait_input\": fantasy_portrait_input, # Fantasy portrait input\n \"phantom_ref\": phantom_ref, # Phantom reference input\n \"reverse_time\": reverse_time, # Reverse RoPE toggle\n \"ntk_alphas\": ntk_alphas, # RoPE freq scaling values\n \"mtv_motion_tokens\": mtv_motion_tokens if mtv_input is not None else None, # MTV-Crafter motion tokens\n \"mtv_motion_rotary_emb\": mtv_motion_rotary_emb if mtv_input is not None else None, # MTV-Crafter RoPE\n \"mtv_strength\": mtv_strength[idx] if mtv_input is not None else 1.0, # MTV-Crafter scaling\n \"mtv_freqs\": mtv_freqs if mtv_input is not None else None, # MTV-Crafter extra RoPE freqs\n \"s2v_audio_input\": s2v_audio_input, # official speech-to-video audio input\n \"s2v_ref_latent\": s2v_ref_latent, # speech-to-video reference latent\n \"s2v_ref_motion\": s2v_ref_motion, # speech-to-video reference motion latent\n \"s2v_audio_scale\": s2v_audio_scale if s2v_audio_input is not None else 1.0, # speech-to-video audio scale\n \"s2v_pose\": s2v_pose if s2v_pose is not None else None, # speech-to-video pose control\n \"s2v_motion_frames\": s2v_motion_frames, # speech-to-video motion frames,\n \"humo_audio\": humo_audio, # humo audio input\n \"humo_audio_scale\": humo_audio_scale if humo_audio is not None else 1,\n \"lynx_embeds\": lynx_embeds, # Lynx face and reference embeddings\n \"x_ovi\": [latent_model_input_ovi.to(z)] if latent_model_input_ovi is not None else None, # Audio latent model input for Ovi\n \"seq_len_ovi\": seq_len_ovi, # Audio latent model sequence length for Ovi\n \"ovi_negative_text_embeds\": ovi_negative_text_embeds, # Audio latent model negative text embeds for Ovi\n \"flashvsr_LQ_latent\": flashvsr_LQ_latent, # FlashVSR LQ latent for upsampling\n \"flashvsr_strength\": flashvsr_strength, # FlashVSR strength\n \"longcat_num_cond_latents\": longcat_num_cond_latents,\n \"longcat_num_ref_latents\": longcat_num_ref_latents,\n \"longcat_avatar_options\": longcat_avatar_options, # LongCat avatar attention options\n \"sdancer_input\": sdancer_input, # SteadyDancer input\n \"one_to_all_input\": one_to_all_data, # One-to-All input\n \"one_to_all_controlnet_strength\": one_to_all_data[\"controlnet_strength\"] if one_to_all_data is not None else 0.0,\n \"scail_input\": scail_data_in, # SCAIL input\n \"dual_control_input\": dual_control_in, # LongVie2 dual control input\n \"context_latents\": context_latents, # Bernini in-context reference\n \"context_window_start\": context_window_start, # Context window start frame index for temporal RoPE alignment\n \"transformer_options\": transformer_options,\n \"rope_negative_offset\": image_embeds.get(\"rope_negative_offset_frames\", 0), # StoryMem rope negative offset\n \"num_memory_frames\": story_mem_latents.shape[1] if story_mem_latents is not None else 0, # StoryMem memory frames\n }\n\n batch_size = 1\n\n if not math.isclose(cfg_scale, 1.0):\n if negative_embeds is None:\n raise ValueError(\"Negative embeddings must be provided for CFG scale > 1.0\")\n if len(positive_embeds) > 1:\n negative_embeds = negative_embeds * len(positive_embeds)\n\n try:\n if not batched_cfg:\n #conditional (positive) pass\n if pos_latent is not None: # for humo\n base_params['x'] = [torch.cat([z[:, :-humo_reference_count], pos_latent], dim=1)]\n base_params[\"add_text_emb\"] = qwenvl_embeds_pos.to(device) if qwenvl_embeds_pos is not None else None # QwenVL embeddings for Bindweave\n noise_pred_cond, noise_pred_ovi, cache_state_cond = transformer(\n context=positive_embeds,\n pred_id=cache_state[0] if cache_state else None,\n vace_data=vace_data, attn_cond=attn_cond,\n **base_params\n )\n noise_pred_cond = noise_pred_cond[0]\n noise_pred_ovi = noise_pred_ovi[0] if noise_pred_ovi is not None else None\n if math.isclose(cfg_scale, 1.0):\n if use_fresca:\n noise_pred_cond = fourier_filter(noise_pred_cond, fresca_scale_low, fresca_scale_high, fresca_freq_cutoff)\n if fantasy_portrait_input is not None and not math.isclose(portrait_cfg[idx], 1.0):\n base_params[\"fantasy_portrait_input\"] = None\n noise_pred_no_portrait, noise_pred_ovi, cache_state_uncond = transformer(context=positive_embeds, pred_id=cache_state[0] if cache_state else None,\n vace_data=vace_data, attn_cond=attn_cond, **base_params)\n return noise_pred_no_portrait[0] + portrait_cfg[idx] * (noise_pred_cond - noise_pred_no_portrait[0]), noise_pred_ovi, [cache_state_cond, cache_state_uncond]\n elif multitalk_audio_input is not None and not math.isclose(audio_cfg_scale[idx], 1.0):\n base_params['multitalk_audio'] = torch.zeros_like(multitalk_audio_input)[-1:]\n noise_pred_uncond_audio, _, cache_state_uncond = transformer(\n context=positive_embeds, pred_id=cache_state[0] if cache_state else None,\n vace_data=vace_data, attn_cond=attn_cond, **base_params)\n return noise_pred_uncond_audio[0] + audio_cfg_scale[idx] * (noise_pred_cond - noise_pred_uncond_audio[0]), noise_pred_ovi, [cache_state_cond, cache_state_uncond]\n else:\n return noise_pred_cond, noise_pred_ovi, [cache_state_cond]\n\n #unconditional (negative) pass\n base_params['is_uncond'] = True\n base_params['clip_fea'] = clip_fea_neg if clip_fea_neg is not None else clip_fea\n base_params[\"add_text_emb\"] = qwenvl_embeds_neg.to(device) if qwenvl_embeds_neg is not None else None # QwenVL embeddings for Bindweave\n base_params['y'] = [image_cond_neg.to(z)] if image_cond_neg is not None else base_params['y']\n if humo_audio_input_neg is not None:\n base_params['humo_audio'] = humo_audio_input_neg\n if neg_latent is not None:\n base_params['x'] = [torch.cat([z[:, :-humo_reference_count], neg_latent], dim=1)]\n\n noise_pred_uncond_text, noise_pred_ovi_uncond, cache_state_uncond = transformer(\n context=negative_embeds if humo_audio_input_neg is None else positive_embeds, #ti #t\n pred_id=cache_state[1] if cache_state else None,\n vace_data=vace_data, attn_cond=attn_cond_neg,\n **base_params)\n noise_pred_uncond_text = noise_pred_uncond_text[0]\n noise_pred_ovi_uncond = noise_pred_ovi_uncond[0] if noise_pred_ovi_uncond is not None else None\n\n # HuMo\n if not math.isclose(humo_audio_cfg_scale[idx], 1.0):\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n if humo_image_cond is not None and humo_audio_input_neg is not None:\n if t > 980 and humo_image_cond_neg_input is not None: # use image cond for first timesteps\n base_params['y'] = [humo_image_cond_neg_input]\n\n noise_pred_humo_audio_uncond, _, cache_state_humo = transformer(\n context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None,\n **base_params)\n\n noise_pred = (noise_pred_uncond_text + humo_audio_cfg_scale[idx] * (noise_pred_cond - noise_pred_humo_audio_uncond[0])\n + (cfg_scale - 2.0) * (noise_pred_humo_audio_uncond[0] - noise_pred_uncond_text))\n return noise_pred, None, [cache_state_cond, cache_state_uncond, cache_state_humo]\n elif humo_audio_input is not None:\n if cache_state is not None and len(cache_state) != 4:\n cache_state.append(None)\n # audio\n noise_pred_humo_null, _, cache_state_humo = transformer(\n context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None,\n **base_params)\n # negative\n if humo_audio_input is not None:\n base_params['humo_audio'] = humo_audio_input\n noise_pred_humo_audio, _, cache_state_humo2 = transformer(\n context=positive_embeds, pred_id=cache_state[3] if cache_state else None, vace_data=None,\n **base_params)\n noise_pred = (humo_audio_cfg_scale[idx] * (noise_pred_cond - noise_pred_humo_audio[0])\n + cfg_scale * (noise_pred_humo_audio[0] - noise_pred_uncond_text)\n + cfg_scale * (noise_pred_uncond_text - noise_pred_humo_null[0])\n + noise_pred_humo_null[0])\n return noise_pred, None, [cache_state_cond, cache_state_uncond, cache_state_humo, cache_state_humo2]\n\n #phantom\n if use_phantom and not math.isclose(phantom_cfg_scale[idx], 1.0):\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n noise_pred_phantom, _, cache_state_phantom = transformer(\n context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None,\n **base_params)\n\n noise_pred = (noise_pred_uncond_text + phantom_cfg_scale[idx] * (noise_pred_phantom[0] - noise_pred_uncond_text)\n + cfg_scale * (noise_pred_cond - noise_pred_phantom[0]))\n return noise_pred, None,[cache_state_cond, cache_state_uncond, cache_state_phantom]\n # audio cfg (fantasytalking and multitalk)\n if (fantasytalking_embeds is not None or multitalk_audio_input is not None):\n if not math.isclose(audio_cfg_scale[idx], 1.0):\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n\n base_params['audio_proj'] = None\n base_params['multitalk_audio'] = torch.zeros_like(multitalk_audio_input)[-1:] if multitalk_audio_input is not None else None\n base_params['is_uncond'] = False\n noise_pred_uncond_audio, _, cache_state_audio = transformer(\n context=negative_embeds,\n pred_id=cache_state[2] if cache_state else None,\n vace_data=vace_data,\n **base_params)\n noise_pred_uncond_audio = noise_pred_uncond_audio[0]\n\n noise_pred = noise_pred_uncond_audio + cfg_scale * (\n (noise_pred_cond - noise_pred_uncond_text)\n + audio_cfg_scale[idx] * (noise_pred_uncond_text - noise_pred_uncond_audio))\n return noise_pred, None,[cache_state_cond, cache_state_uncond, cache_state_audio]\n # lynx\n if lynx_embeds is not None and not math.isclose(lynx_cfg_scale[idx], 1.0):\n base_params['is_uncond'] = False\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n noise_pred_lynx, _, cache_state_lynx = transformer(\n context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None,\n **base_params)\n\n noise_pred = (noise_pred_uncond_text + lynx_cfg_scale[idx] * (noise_pred_lynx[0] - noise_pred_uncond_text)\n + cfg_scale * (noise_pred_cond - noise_pred_lynx[0]))\n return noise_pred, None, [cache_state_cond, cache_state_uncond, cache_state_lynx]\n # one-to-all\n if one_to_all_data is not None and not math.isclose(one_to_all_pose_cfg_scale[idx], 1.0):\n tqdm.write(\"One-to-All pose CFG pass...\")\n base_params['is_uncond'] = False\n base_params['one_to_all_controlnet_strength'] = 0.0\n if cache_state is not None and len(cache_state) != 3:\n cache_state.append(None)\n noise_pred_pose_uncond, _, cache_state_ref = transformer(\n context=negative_embeds, pred_id=cache_state[2] if cache_state else None, vace_data=None,\n **base_params)\n\n noise_pred = (noise_pred_uncond_text + one_to_all_pose_cfg_scale[idx] * (noise_pred_pose_uncond[0] - noise_pred_uncond_text)\n + cfg_scale * (noise_pred_cond - noise_pred_pose_uncond[0]))\n return noise_pred, None, [cache_state_cond, cache_state_uncond, cache_state_ref]\n\n #batched\n else:\n base_params[\"x\"] = [z, z]\n base_params[\"y\"] = [image_cond_input, image_cond_input] if image_cond_input is not None else None\n base_params['clip_fea'] = torch.cat([clip_fea, clip_fea], dim=0)\n cache_state_uncond = None\n [noise_pred_cond, noise_pred_uncond_text], _, cache_state_cond = transformer(\n context=positive_embeds + negative_embeds, is_uncond=False,\n pred_id=cache_state[0] if cache_state else None,\n **base_params\n )\n except Exception as e:\n log.error(f\"Error during model prediction: {e}\")\n if force_offload:\n if not model[\"auto_cpu_offload\"]:\n offload_transformer(transformer)\n raise e\n\n #https://github.com/WeichenFan/CFG-Zero-star/\n alpha = 1.0\n if use_cfg_zero_star:\n alpha = optimized_scale(\n noise_pred_cond.view(batch_size, -1),\n noise_pred_uncond_text.view(batch_size, -1)\n ).view(batch_size, 1, 1, 1)\n\n noise_pred_uncond_text = noise_pred_uncond_text * alpha\n\n if use_tangential:\n noise_pred_uncond_text = tangential_projection(noise_pred_cond, noise_pred_uncond_text)\n\n # RAAG (RATIO-aware Adaptive Guidance)\n if raag_alpha > 0.0:\n cfg_scale = get_raag_guidance(noise_pred_cond, noise_pred_uncond_text, cfg_scale, raag_alpha)\n log.info(f\"RAAG modified cfg: {cfg_scale}\")\n\n #https://github.com/WikiChao/FreSca\n if use_fresca:\n filtered_cond = fourier_filter(noise_pred_cond - noise_pred_uncond_text, fresca_scale_low, fresca_scale_high, fresca_freq_cutoff)\n noise_pred = noise_pred_uncond_text + cfg_scale * filtered_cond * alpha\n elif uses_apg_guidance(guidance_mode):\n sigma = sample_scheduler.sigmas[idx]\n x_cond = z - sigma * noise_pred_cond\n x_uncond = z - sigma * noise_pred_uncond_text\n x_guided = _normalized_guidance(x_cond, x_uncond, cfg_scale, _momentum_buf, apg_eta, apg_norm_threshold)\n noise_pred = (z - x_guided) / sigma\n else:\n noise_pred = noise_pred_uncond_text + cfg_scale * (noise_pred_cond - noise_pred_uncond_text)\n del noise_pred_uncond_text, noise_pred_cond\n\n if latent_model_input_ovi is not None:\n if ovi_audio_cfg is None:\n audio_cfg_scale = cfg_scale - 1.0 if cfg_scale > 4.0 else cfg_scale\n else:\n audio_cfg_scale = ovi_audio_cfg[idx]\n noise_pred_ovi = noise_pred_ovi_uncond + audio_cfg_scale * (noise_pred_ovi - noise_pred_ovi_uncond)\n\n return noise_pred, noise_pred_ovi, [cache_state_cond, cache_state_uncond]\n\n if args.preview_method in [LatentPreviewMethod.Auto, LatentPreviewMethod.Latent2RGB]: #default for latent2rgb\n from latent_preview import prepare_callback\n else:\n from .latent_preview import prepare_callback #custom for tiny VAE previews\n callback = prepare_callback(patcher, len(timesteps))\n callback_latent = None\n\n if not multitalk_sampling and not framepack:\n log.info(\"-\" * 10 + \" Sampling start \" + \"-\" * 10)\n log.info(f\"{(latent_video_length-1) * 4 + 1} frames at {latent.shape[3]*vae_upscale_factor}x{latent.shape[2]*vae_upscale_factor} (Input sequence length: {seq_len}) with {steps-ttm_start_step} steps\")\n\n\n # Differential diffusion prep\n masks = None\n if not multitalk_sampling and samples is not None and noise_mask is not None:\n thresholds = torch.arange(len(timesteps), dtype=original_image.dtype) / len(timesteps)\n thresholds = thresholds.reshape(-1, 1, 1, 1, 1).to(device)\n masks = (1-noise_mask.repeat(len(timesteps), 1, 1, 1, 1).to(device)) > thresholds\n\n latent_shift_loop = False\n if loop_args is not None:\n latent_shift_loop = is_looped = True\n latent_skip = loop_args[\"shift_skip\"]\n latent_shift_start_percent = loop_args[\"start_percent\"]\n latent_shift_end_percent = loop_args[\"end_percent\"]\n shift_idx = 0\n\n #clear memory before sampling\n mm.soft_empty_cache()\n gc.collect()\n try:\n torch.cuda.reset_peak_memory_stats(device)\n except Exception:\n pass\n\n # Bernini runtime: dedicated denoise loop + slim transformer forward\n from .bernini_runtime import bernini_runtime_eligible\n from .bernini_denoise import BerniniDenoiseConfig, run_bernini_denoise\n\n _use_bernini_runtime = (\n bernini_runtime_eligible(\n image_embeds,\n context_options=context_options,\n transformer=transformer,\n unianimate_poses=unianimate_poses,\n fantasytalking_embeds=fantasytalking_embeds,\n uni3c_embeds=uni3c_embeds,\n multitalk_embeds=multitalk_embeds,\n loop_args=loop_args,\n experimental_args=experimental_args,\n freeinit_args=freeinit_args,\n flowedit_args=flowedit_args,\n feta_args=feta_args,\n batched_cfg=batched_cfg,\n )\n and not multitalk_sampling\n and not framepack\n and not framepack\n and freeinit_args is None\n )\n\n if _use_bernini_runtime:\n from .bernini_perf import auto_batched_cfg\n\n effective_batched_cfg = batched_cfg or auto_batched_cfg(text_embeds, cfg, image_embeds)\n if effective_batched_cfg and not batched_cfg:\n log.info(\"Bernini perf: batched CFG auto-enabled\")\n context_fn = None\n if context_options is not None:\n from .context_windows.context import get_context_scheduler\n context_fn = get_context_scheduler(context_options[\"context_schedule\"])\n denoise_cfg = BerniniDenoiseConfig(\n transformer=transformer,\n patcher=patcher,\n model_meta=model,\n sample_scheduler=sample_scheduler,\n timesteps=timesteps,\n cfg=cfg,\n text_embeds=text_embeds,\n seq_len=seq_len,\n freqs=freqs,\n context_latents=context_latents,\n latent_video_length=latent_video_length,\n steps=steps,\n ttm_start_step=ttm_start_step,\n device=device,\n dtype=dtype,\n guidance_mode=guidance_mode,\n apg_eta=apg_eta,\n apg_momentum=apg_momentum,\n apg_norm_threshold=apg_norm_threshold,\n context_options=context_options,\n cache_state=self.cache_state,\n force_offload=force_offload,\n seed_g=seed_g,\n scheduler_step_args=scheduler_step_args,\n transformer_options=transformer_options or {},\n vae_upscale_factor=vae_upscale_factor,\n is_looped=is_looped,\n context_fn=context_fn,\n batched_cfg=effective_batched_cfg,\n masks=masks,\n original_image=original_image.squeeze(0) if original_image is not None else None,\n noise=noise,\n )\n latent, callback_latent = run_bernini_denoise(denoise_cfg, latent, callback=callback)\n\n # Main sampling loop with FreeInit iterations (skipped when Bernini runtime ran)\n iterations = 0 if _use_bernini_runtime else (freeinit_args.get(\"freeinit_num_iters\", 3) if freeinit_args is not None else 1)\n current_latent = latent\n initial_noise_saved = None\n\n for iter_idx in range(iterations):\n\n # FreeInit noise reinitialization (after first iteration)\n if freeinit_args is not None and iter_idx > 0:\n # restart scheduler for each iteration\n sample_scheduler, timesteps,_,_ = get_scheduler(scheduler, steps, start_step, end_step, shift, device, transformer.dim, denoise_strength, sigmas=sigmas)\n\n # Re-apply start_step and end_step logic to timesteps and sigmas\n if end_step != -1:\n timesteps = timesteps[:end_step]\n sample_scheduler.sigmas = sample_scheduler.sigmas[:end_step+1]\n if start_step > 0:\n timesteps = timesteps[start_step:]\n sample_scheduler.sigmas = sample_scheduler.sigmas[start_step:]\n if hasattr(sample_scheduler, 'timesteps'):\n sample_scheduler.timesteps = timesteps\n\n # Diffuse current latent to t=999\n diffuse_timesteps = torch.full((noise.shape[0],), 999, device=device, dtype=torch.long)\n z_T = add_noise(\n current_latent.to(device),\n initial_noise_saved.to(device),\n diffuse_timesteps\n )\n\n # Generate new random noise\n z_rand = torch.randn(z_T.shape, dtype=torch.float32, generator=seed_g, device=torch.device(\"cpu\"))\n # Apply frequency mixing\n current_latent = (freq_mix_3d(z_T.to(torch.float32), z_rand.to(device), LPF=freq_filter)).to(dtype)\n\n # Store initial noise for first iteration\n if freeinit_args is not None and iter_idx == 0:\n initial_noise_saved = current_latent.detach().clone()\n if input_samples is not None:\n current_latent = input_samples.to(device)\n continue\n\n # Reset per-iteration states\n self.cache_state = [None, None]\n self.cache_state_source = [None, None]\n self.cache_states_context = []\n if context_options is not None:\n self.window_tracker = WindowTracker(verbose=context_options[\"verbose\"])\n\n # Set latent for denoising\n latent = current_latent\n\n if is_pusa and clean_latent_indices:\n pusa_noisy_steps = image_embeds.get(\"pusa_noisy_steps\", -1)\n if pusa_noisy_steps == -1:\n pusa_noisy_steps = len(timesteps)\n try:\n pbar = ProgressBar(len(timesteps) - ttm_start_step)\n #region main loop start\n for idx, t in enumerate(tqdm(timesteps[ttm_start_step:], disable=multitalk_sampling)):\n\n if bidirectional_sampling:\n latent_flipped = torch.flip(latent, dims=[1])\n latent_model_input_flipped = latent_flipped.to(device)\n\n self.noise_front_pad_num = 0\n\n #InfiniteTalk first frame handling\n if (extra_latents is not None\n and not multitalk_sampling\n and transformer.multitalk_model_type==\"InfiniteTalk\"):\n for entry in extra_latents:\n add_index = entry[\"index\"]\n num_extra_frames = entry[\"samples\"].shape[2]\n latent[:, add_index:add_index+num_extra_frames] = entry[\"samples\"].to(latent)\n\n latent_model_input = latent.to(device)\n latent_model_input_ovi = latent_ovi.to(device) if latent_ovi is not None else None\n\n current_step_percentage = idx / len(timesteps)\n\n timestep = torch.tensor([t]).to(device)\n if is_pusa or ((is_5b or transformer.is_longcat) and clean_latent_indices):\n orig_timestep = timestep\n timestep = timestep.unsqueeze(1).repeat(1, latent_video_length)\n if extra_latents is not None:\n if clean_latent_indices and noise_multipliers is not None:\n if is_pusa:\n scheduler_step_args[\"cond_frame_latent_indices\"] = clean_latent_indices\n scheduler_step_args[\"noise_multipliers\"] = noise_multipliers\n for latent_idx in clean_latent_indices:\n timestep[:, latent_idx] = timestep[:, latent_idx] * noise_multipliers[latent_idx]\n # add noise for conditioning frames if multiplier > 0\n if idx < pusa_noisy_steps and noise_multipliers[latent_idx] > 0:\n latent_size = (1, latent.shape[0], latent.shape[1], latent.shape[2], latent.shape[3])\n noise_for_cond = torch.randn(latent_size, generator=seed_g, device=torch.device(\"cpu\"))\n timestep_cond = torch.ones_like(timestep) * timestep.max()\n if is_pusa:\n latent[:, latent_idx:latent_idx+1] = sample_scheduler.add_noise_for_conditioning_frames(\n latent[:, latent_idx:latent_idx+1].to(device),\n noise_for_cond[:, :, latent_idx:latent_idx+1].to(device),\n timestep_cond[:, latent_idx:latent_idx+1].to(device),\n noise_multiplier=noise_multipliers[latent_idx])\n else:\n timestep[:, clean_latent_indices] = 0\n #print(\"timestep: \", timestep)\n\n ### latent shift\n if latent_shift_loop:\n if latent_shift_start_percent <= current_step_percentage <= latent_shift_end_percent:\n latent_model_input = torch.cat([latent_model_input[:, shift_idx:]] + [latent_model_input[:, :shift_idx]], dim=1)\n\n #enhance-a-video\n enhance_enabled = False\n if feta_args is not None and feta_start_percent <= current_step_percentage <= feta_end_percent:\n enhance_enabled = True\n #region context windowing\n if context_options is not None:\n counter = torch.zeros_like(latent_model_input, device=device)\n noise_pred = torch.zeros_like(latent_model_input, device=device)\n if image_cond is not None and image_cond.ndim > 1:\n latent_video_length = min(latent_video_length, image_cond.shape[1])\n context_queue = list(context(idx, steps, latent_video_length, context_frames, context_stride, context_overlap))\n fraction_per_context = 1.0 / len(context_queue)\n context_pbar = ProgressBar(steps)\n step_start_progress = idx\n\n # Validate all context windows before processing\n max_idx = latent_model_input.shape[1] if latent_model_input.ndim > 1 else 0\n for window_indices in context_queue:\n if not all(0 <= idx < max_idx for idx in window_indices):\n raise ValueError(f\"Invalid context window indices {window_indices} for latent_model_input with shape {latent_model_input.shape}\")\n\n for i, c in enumerate(context_queue):\n window_id = self.window_tracker.get_window_id(c)\n\n if cache_args is not None:\n current_teacache = self.window_tracker.get_teacache(window_id, self.cache_state)\n else:\n current_teacache = None\n\n prompt_index = min(int(max(c) / section_size), num_prompts - 1)\n if context_options[\"verbose\"]:\n log.info(f\"Prompt index: {prompt_index}\")\n\n # Use the appropriate prompt for this section\n if len(text_embeds[\"prompt_embeds\"]) > 1:\n positive = [text_embeds[\"prompt_embeds\"][prompt_index]]\n else:\n positive = text_embeds[\"prompt_embeds\"]\n\n partial_img_emb = partial_control_latents = None\n if image_cond is not None:\n partial_img_emb = image_cond[:, c].to(device)\n if c[0] != 0 and context_reference_latent is not None:\n if context_reference_latent.shape[0] == 1: #only single extra init latent\n new_init_image = context_reference_latent[0, :, 0].to(device)\n partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0)\n elif context_reference_latent.shape[0] > 1:\n num_extra_inits = context_reference_latent.shape[0]\n section_size = (latent_video_length / num_extra_inits)\n extra_init_index = min(int(max(c) / section_size), num_extra_inits - 1)\n if context_options[\"verbose\"]:\n log.info(f\"extra init image index: {extra_init_index}\")\n new_init_image = context_reference_latent[extra_init_index, :, 0].to(device)\n partial_img_emb[:, 0] = torch.cat([image_cond[:4, 0].to(device), new_init_image], dim=0)\n else:\n new_init_image = image_cond[:, 0].to(device)\n partial_img_emb[:, 0] = new_init_image\n\n if control_latents is not None:\n partial_control_latents = control_latents[:, c]\n\n partial_control_camera_latents = None\n if control_camera_latents is not None:\n partial_control_camera_latents = control_camera_latents[:, :, c]\n\n partial_vace_context = None\n if vace_data is not None:\n window_vace_data = []\n for vace_entry in vace_data:\n partial_context = vace_entry[\"context\"][0][:, c]\n if has_ref:\n if c[0] != 0 and context_reference_latent is not None:\n if context_reference_latent.shape[0] == 1: #only single extra init latent\n partial_context[16:32, :1] = context_reference_latent[0, :, :1].to(device)\n elif context_reference_latent.shape[0] > 1:\n num_extra_inits = context_reference_latent.shape[0]\n section_size = (latent_video_length / num_extra_inits)\n extra_init_index = min(int(max(c) / section_size), num_extra_inits - 1)\n if context_options[\"verbose\"]:\n log.info(f\"extra init image index: {extra_init_index}\")\n partial_context[16:32, :1] = context_reference_latent[extra_init_index, :, :1].to(device)\n else:\n partial_context[:, 0] = vace_entry[\"context\"][0][:, 0]\n\n window_vace_data.append({\n \"context\": [partial_context],\n \"scale\": vace_entry[\"scale\"],\n \"start\": vace_entry[\"start\"],\n \"end\": vace_entry[\"end\"],\n \"seq_len\": vace_entry[\"seq_len\"]\n })\n\n partial_vace_context = window_vace_data\n\n partial_audio_proj = None\n if fantasytalking_embeds is not None:\n partial_audio_proj = audio_proj[:, c]\n\n partial_fantasy_portrait_input = None\n if fantasy_portrait_input is not None:\n partial_fantasy_portrait_input = fantasy_portrait_input.copy()\n partial_fantasy_portrait_input[\"adapter_proj\"] = fantasy_portrait_input[\"adapter_proj\"][:, c]\n\n partial_latent_model_input = latent_model_input[:, c]\n if latents_to_insert is not None and c[0] != 0:\n partial_latent_model_input[:, :1] = latents_to_insert\n\n partial_unianim_data = None\n if unianim_data is not None:\n partial_dwpose = dwpose_data[:, :, c]\n partial_unianim_data = {\n \"dwpose\": partial_dwpose,\n \"random_ref\": unianim_data[\"random_ref\"],\n \"strength\": unianimate_poses[\"strength\"],\n \"start_percent\": unianimate_poses[\"start_percent\"],\n \"end_percent\": unianimate_poses[\"end_percent\"]\n }\n\n partial_mtv_motion_tokens = None\n if mtv_input is not None:\n start_token_index = c[0] * 24\n end_token_index = (c[-1] + 1) * 24\n partial_mtv_motion_tokens = mtv_motion_tokens[:, start_token_index:end_token_index, :]\n if context_options[\"verbose\"]:\n log.info(f\"context window: {c}\")\n log.info(f\"motion_token_indices: {start_token_index}-{end_token_index}\")\n\n partial_s2v_audio_input = None\n if s2v_audio_input is not None:\n audio_start = c[0] * 4\n audio_end = c[-1] * 4 + 1\n center_indices = torch.arange(audio_start, audio_end, 1)\n center_indices = torch.clamp(center_indices, min=0, max=s2v_audio_input.shape[-1] - 1)\n partial_s2v_audio_input = s2v_audio_input[..., center_indices]\n\n partial_s2v_pose = None\n if s2v_pose is not None:\n partial_s2v_pose = s2v_pose[:, :, c].to(device, dtype)\n\n partial_add_cond = None\n if add_cond is not None:\n partial_add_cond = add_cond[:, :, c].to(device, dtype)\n\n partial_flashvsr_LQ_latent = None\n if LQ_images is not None:\n start = c[0] * 4\n end = c[-1] * 4 + 1 + 4\n center_indices = torch.arange(start, end, 1)\n center_indices = torch.clamp(center_indices, min=0, max=LQ_images.shape[2] - 1)\n partial_flashvsr_LQ_images = LQ_images[:, :, center_indices].to(device)\n partial_flashvsr_LQ_latent = transformer.LQ_proj_in(partial_flashvsr_LQ_images)\n\n if len(timestep.shape) != 1:\n partial_timestep = timestep[:, c]\n partial_timestep[:, :1] = 0\n else:\n partial_timestep = timestep\n\n orig_model_input_frames = partial_latent_model_input.shape[1]\n\n sliced_context_latents = None\n if context_latents is not None:\n sliced_context_latents = []\n for lat in context_latents:\n if lat.shape[1] > 1 and lat.shape[1] == noise.shape[1]:\n sliced_context_latents.append(lat[:, c].to(device))\n else:\n sliced_context_latents.append(lat.to(device))\n noise_pred_context, _, new_teacache = predict_with_cfg(\n partial_latent_model_input,\n cfg[idx], positive,\n text_embeds[\"negative_prompt_embeds\"],\n partial_timestep, idx, partial_img_emb, clip_fea, partial_control_latents, partial_vace_context, partial_unianim_data,partial_audio_proj,\n partial_control_camera_latents, partial_add_cond, current_teacache, context_window=c, fantasy_portrait_input=partial_fantasy_portrait_input,\n mtv_motion_tokens=partial_mtv_motion_tokens, s2v_audio_input=partial_s2v_audio_input, s2v_motion_frames=[1, 0], s2v_pose=partial_s2v_pose,\n humo_image_cond=humo_image_cond, humo_image_cond_neg=humo_image_cond_neg, humo_audio=humo_audio, humo_audio_neg=humo_audio_neg,\n multitalk_audio_embeds=multitalk_audio_embeds,\n uni3c_data=uni3c_data, flashvsr_LQ_latent=partial_flashvsr_LQ_latent, context_latents=sliced_context_latents,\n context_window_start=c[0])\n\n if cache_args is not None:\n self.window_tracker.cache_states[window_id] = new_teacache\n\n if mocha_embeds is not None:\n noise_pred_context = noise_pred_context[:, :orig_model_input_frames]\n\n window_mask = create_window_mask(noise_pred_context, c, noise.shape[1], context_overlap, looped=is_looped, window_type=context_options[\"fuse_method\"])\n noise_pred[:, c] += noise_pred_context * window_mask\n counter[:, c] += window_mask\n context_pbar.update_absolute(step_start_progress + (i + 1) * fraction_per_context, len(timesteps))\n noise_pred /= counter\n # region framepack loop\n elif framepack:\n framepack_out = []\n ref_motion_image = None\n #infer_frames = image_embeds[\"num_frames\"]\n infer_frames = s2v_audio_embeds.get(\"frame_window_size\", 80)\n motion_frames = infer_frames - 7 #73 default\n lat_motion_frames = (motion_frames + 3) // 4\n lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames\n\n step_iteration_count = 0\n total_frames = s2v_audio_input.shape[-1]\n\n s2v_motion_frames = [motion_frames, lat_motion_frames]\n\n noise = torch.randn( #C, T, H, W\n 48 if is_5b else 16,\n lat_target_frames,\n target_shape[2],\n target_shape[3],\n dtype=torch.float32,\n generator=seed_g,\n device=torch.device(\"cpu\"))\n\n seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * noise.shape[1])\n\n if ref_motion_image is None:\n ref_motion_image = torch.zeros(\n [1, 3, motion_frames, latent.shape[2]*vae_upscale_factor, latent.shape[3]*vae_upscale_factor],\n dtype=vae.dtype,\n device=device)\n videos_last_frames = ref_motion_image\n\n if s2v_pose is not None:\n pose_cond_list = []\n for r in range(s2v_num_repeat):\n pose_start = r * (infer_frames // 4)\n pose_end = pose_start + (infer_frames // 4)\n\n cond_lat = s2v_pose[:, :, pose_start:pose_end]\n\n pad_len = (infer_frames // 4) - cond_lat.shape[2]\n if pad_len > 0:\n pad = -torch.ones(cond_lat.shape[0], cond_lat.shape[1], pad_len, cond_lat.shape[3], cond_lat.shape[4], device=cond_lat.device, dtype=cond_lat.dtype)\n cond_lat = torch.cat([cond_lat, pad], dim=2)\n pose_cond_list.append(cond_lat.cpu())\n\n log.info(f\"Sampling {total_frames} frames in {s2v_num_repeat} windows, at {latent.shape[3]*vae_upscale_factor}x{latent.shape[2]*vae_upscale_factor} with {steps} steps\")\n # sample\n for r in range(s2v_num_repeat):\n\n mm.soft_empty_cache()\n gc.collect()\n if ref_motion_image is not None:\n vae.to(device)\n ref_motion = vae.encode(ref_motion_image.to(vae.dtype), device=device, pbar=False).to(dtype)[0]\n\n vae.to(offload_device)\n\n left_idx = r * infer_frames\n right_idx = r * infer_frames + infer_frames\n\n s2v_audio_input_slice = s2v_audio_input[..., left_idx:right_idx]\n if s2v_audio_input_slice.shape[-1] < (right_idx - left_idx):\n pad_len = (right_idx - left_idx) - s2v_audio_input_slice.shape[-1]\n pad_shape = list(s2v_audio_input_slice.shape)\n pad_shape[-1] = pad_len\n pad = torch.zeros(pad_shape, device=s2v_audio_input_slice.device, dtype=s2v_audio_input_slice.dtype)\n log.info(f\"Padding s2v_audio_input_slice from {s2v_audio_input_slice.shape[-1]} to {right_idx - left_idx}\")\n s2v_audio_input_slice = torch.cat([s2v_audio_input_slice, pad], dim=-1)\n\n if ref_motion_image is not None:\n input_motion_latents = ref_motion.clone().unsqueeze(0)\n else:\n input_motion_latents = None\n\n s2v_pose_slice = None\n if s2v_pose is not None:\n s2v_pose_slice = pose_cond_list[r].to(device)\n\n if isinstance(scheduler, dict):\n sample_scheduler = copy.deepcopy(scheduler[\"sample_scheduler\"])\n timesteps = scheduler[\"timesteps\"]\n else:\n sample_scheduler, timesteps,_,_ = get_scheduler(scheduler, total_steps, start_step, end_step, shift, device, transformer.dim, denoise_strength, sigmas=sigmas)\n\n latent = noise.to(device)\n for i, t in enumerate(tqdm(timesteps, desc=f\"Sampling audio indices {left_idx}-{right_idx}\", position=0)):\n latent_model_input = latent.to(device)\n timestep = torch.tensor([t]).to(device)\n noise_pred, _, self.cache_state = predict_with_cfg(\n latent_model_input,\n cfg[idx],\n text_embeds[\"prompt_embeds\"],\n text_embeds[\"negative_prompt_embeds\"],\n timestep, idx, image_cond, clip_fea, control_latents, vace_data, unianim_data, audio_proj, control_camera_latents, add_cond,\n cache_state=self.cache_state, fantasy_portrait_input=fantasy_portrait_input, mtv_motion_tokens=mtv_motion_tokens,\n context_latents=context_latents,\n s2v_audio_input=s2v_audio_input_slice, s2v_ref_motion=input_motion_latents, s2v_motion_frames=s2v_motion_frames, s2v_pose=s2v_pose_slice)\n\n latent = sample_scheduler.step(\n noise_pred.unsqueeze(0), timestep, latent.unsqueeze(0),\n **scheduler_step_args)[0].squeeze(0).detach()\n if callback is not None:\n callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device) / 1000).detach().permute(1,0,2,3)\n callback(step_iteration_count, callback_latent, None, s2v_num_repeat*(len(timesteps)))\n del callback_latent\n step_iteration_count += 1\n del latent_model_input, noise_pred\n\n\n vae.to(device)\n decode_latents = torch.cat([ref_motion.unsqueeze(0), latent.unsqueeze(0)], dim=2)\n image = vae.decode(decode_latents.to(device, vae.dtype), device=device, pbar=False)[0]\n del decode_latents\n image = image.unsqueeze(0)[:, :, -infer_frames:]\n if r == 0:\n image = image[:, :, 3:]\n\n framepack_out.append(image.cpu())\n\n overlap_frames_num = min(motion_frames, image.shape[2])\n\n videos_last_frames = torch.cat([\n videos_last_frames[:, :, overlap_frames_num:],\n image[:, :, -overlap_frames_num:]], dim=2).to(device, vae.dtype)\n\n ref_motion_image = videos_last_frames\n\n vae.to(offload_device)\n\n mm.soft_empty_cache()\n gen_video_samples = torch.cat(framepack_out, dim=2).squeeze(0).permute(1, 2, 3, 0)\n\n if force_offload:\n if not model[\"auto_cpu_offload\"]:\n offload_transformer(transformer)\n try:\n print_memory(device)\n torch.cuda.reset_peak_memory_stats(device)\n except Exception:\n pass\n return {\"video\": gen_video_samples},\n #region normal inference\n else:\n noise_pred, noise_pred_ovi, self.cache_state = predict_with_cfg(\n latent_model_input,\n cfg[idx], text_embeds[\"prompt_embeds\"], text_embeds[\"negative_prompt_embeds\"],\n timestep, idx, image_cond, clip_fea, control_latents, vace_data, unianim_data, audio_proj, control_camera_latents, add_cond,\n cache_state=self.cache_state, fantasy_portrait_input=fantasy_portrait_input, multitalk_audio_embeds=multitalk_audio_embeds, mtv_motion_tokens=mtv_motion_tokens, s2v_audio_input=s2v_audio_input,\n humo_image_cond=humo_image_cond, humo_image_cond_neg=humo_image_cond_neg, humo_audio=humo_audio, humo_audio_neg=humo_audio_neg,\n uni3c_data = uni3c_data, latent_model_input_ovi=latent_model_input_ovi, flashvsr_LQ_latent=flashvsr_LQ_latent,\n context_latents=context_latents,\n )\n if bidirectional_sampling:\n noise_pred_flipped, _,self.cache_state = predict_with_cfg(\n latent_model_input_flipped,\n cfg[idx], text_embeds[\"prompt_embeds\"], text_embeds[\"negative_prompt_embeds\"],\n timestep, idx, image_cond, clip_fea, control_latents, vace_data, unianim_data, audio_proj, control_camera_latents, add_cond,\n cache_state=self.cache_state, fantasy_portrait_input=fantasy_portrait_input, mtv_motion_tokens=mtv_motion_tokens,reverse_time=True,\n context_latents=context_latents,)\n\n if latent_shift_loop:\n #reverse latent shift\n if latent_shift_start_percent <= current_step_percentage <= latent_shift_end_percent:\n noise_pred = torch.cat([noise_pred[:, latent_video_length - shift_idx:]] + [noise_pred[:, :latent_video_length - shift_idx]], dim=1)\n shift_idx = (shift_idx + latent_skip) % latent_video_length\n\n latent = latent.to(device)\n\n if self.noise_front_pad_num > 0:\n noise_pred = noise_pred[:, self.noise_front_pad_num:]\n\n if use_tsr:\n noise_pred = temporal_score_rescaling(noise_pred, latent, timestep, tsr_k, tsr_sigma)\n\n if transformer.is_longcat:\n noise_pred = -noise_pred\n\n if len(timestep.shape) != 1 and clean_latent_indices and not is_pusa: #5b and longcat, skip clean latents for scheduler step\n step_process_indices = [i for i in range(latent.shape[1]) if i not in clean_latent_indices]\n latent[:, step_process_indices] = sample_scheduler.step(noise_pred[:, step_process_indices].unsqueeze(0), orig_timestep,\n latent[:, step_process_indices].unsqueeze(0), **scheduler_step_args)[0].squeeze(0).detach()\n else:\n if latents_to_not_step > 0:\n raw_latent = latent[:, :latents_to_not_step]\n noise_pred_in = noise_pred[:, latents_to_not_step:]\n latent = latent[:, latents_to_not_step:]\n elif recammaster is not None or mocha_embeds is not None:\n noise_pred_in = noise_pred[:, :orig_noise_len]\n latent = latent[:, :orig_noise_len]\n else:\n noise_pred_in = noise_pred\n latent = sample_scheduler.step(noise_pred_in.unsqueeze(0), timestep, latent.unsqueeze(0), **scheduler_step_args)[0].squeeze(0).detach()\n if noise_pred_flipped is not None:\n latent_backwards = sample_scheduler_flipped.step(noise_pred_flipped.unsqueeze(0), timestep, latent_flipped.unsqueeze(0), **scheduler_step_args)[0].squeeze(0).detach()\n latent_backwards = torch.flip(latent_backwards, dims=[1])\n latent = latent * 0.5 + latent_backwards * 0.5\n if latents_to_not_step > 0:\n latent = torch.cat([raw_latent, latent], dim=1)\n\n latent = latent.detach()\n\n if latent_ovi is not None:\n latent_ovi = sample_scheduler_ovi.step(noise_pred_ovi.unsqueeze(0), t, latent_ovi.to(device).unsqueeze(0), **scheduler_step_args)[0].squeeze(0).detach()\n\n #InfiniteTalk first frame handling\n if (extra_latents is not None\n and not multitalk_sampling\n and transformer.multitalk_model_type==\"InfiniteTalk\"):\n for entry in extra_latents:\n add_index = entry[\"index\"]\n num_extra_frames = entry[\"samples\"].shape[2]\n latent[:, add_index:add_index+num_extra_frames] = entry[\"samples\"].to(latent)\n\n # differential diffusion inpaint\n if masks is not None:\n if idx < len(timesteps) - 1:\n noise_timestep = timesteps[idx+1]\n image_latent = sample_scheduler.scale_noise(\n original_image.to(device), torch.tensor([noise_timestep]), noise.to(device)\n )\n mask = masks[idx].to(latent)\n latent = image_latent * mask + latent * (1-mask)\n\n # TTM\n if ttm_reference_latents is not None and (idx + ttm_start_step) < ttm_end_step:\n if idx + ttm_start_step + 1 < len(sample_scheduler.all_timesteps):\n noisy_latents = add_noise(ttm_reference_latents, noise, sample_scheduler.all_timesteps[idx + ttm_start_step + 1].to(noise.device)).to(latent)\n latent = latent * (1 - motion_mask) + noisy_latents * motion_mask\n else:\n latent = latent * (1 - motion_mask) + ttm_reference_latents.to(latent) * motion_mask\n\n if freeinit_args is not None:\n current_latent = latent.clone()\n\n if callback is not None:\n if recammaster is not None or mocha_embeds is not None:\n callback_latent = (latent_model_input[:, :orig_noise_len].to(device) - noise_pred[:, :orig_noise_len].to(device) * t.to(device) / 1000).detach()\n #elif phantom_latents is not None:\n # callback_latent = (latent_model_input[:,:-phantom_latents.shape[1]].to(device) - noise_pred[:,:-phantom_latents.shape[1]].to(device) * t.to(device) / 1000).detach()\n elif humo_reference_count > 0:\n callback_latent = (latent_model_input[:,:-humo_reference_count].to(device) - noise_pred[:,:-humo_reference_count].to(device) * t.to(device) / 1000).detach()\n elif \"rcm\" in sample_scheduler.__class__.__name__.lower():\n callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device)).detach()\n else:\n callback_latent = (latent_model_input.to(device) - noise_pred.to(device) * t.to(device) / 1000).detach()\n callback(idx, callback_latent.permute(1,0,2,3), None, len(timesteps))\n else:\n pbar.update(1)\n\n noise_pred = noise_pred_in = noise_pred_ovi = None\n noise_pred_flipped = None\n latent_model_input = latent_model_input_ovi = None\n latent_model_input_flipped = None\n latent_flipped = latent_backwards = raw_latent = None\n timestep = orig_timestep = None\n\n except Exception as e:\n log.error(f\"Error during sampling: {e}\")\n raise\n finally:\n if force_offload and not model[\"auto_cpu_offload\"]:\n offload_transformer(transformer)\n\n if phantom_latents is not None:\n latent = latent[:,:-phantom_latents.shape[1]]\n if humo_reference_count > 0:\n latent = latent[:,:-humo_reference_count]\n if longcat_ref_latent is not None:\n latent = latent[:, longcat_ref_latent.shape[1]:]\n if story_mem_latents is not None:\n latent = latent[:, story_mem_latents.shape[1]:]\n\n log.info(\"-\" * 10 + \" Sampling end \" + \"-\" * 12)\n\n cache_states = None\n if cache_args is not None:\n cache_report(transformer, cache_args)\n if end_step != -1 and end_step < total_steps:\n cache_states = {\n \"cache_state\": self.cache_state,\n \"easycache_state\": transformer.easycache_state,\n \"teacache_state\": transformer.teacache_state,\n \"magcache_state\": transformer.magcache_state,\n }\n\n try:\n print_memory(device)\n torch.cuda.reset_peak_memory_stats(device)\n except Exception:\n pass\n return ({\n \"samples\": latent.unsqueeze(0).cpu(),\n \"looped\": is_looped,\n \"end_image\": end_image if not fun_or_fl2v_model else None,\n \"has_ref\": has_ref,\n \"drop_last\": drop_last,\n \"generator_state\": seed_g.get_state(),\n \"original_image\": original_image.cpu() if original_image is not None else None,\n \"noise_mask\": input_noise_mask.cpu() if input_noise_mask is not None else None,\n \"cache_states\": cache_states,\n \"latent_ovi_audio\": latent_ovi.unsqueeze(0).transpose(1, 2).cpu() if latent_ovi is not None else None,\n \"flashvsr_LQ_images\": LQ_images,\n },{\n \"samples\": callback_latent.unsqueeze(0).cpu() if callback_latent is not None else None,\n })\n\nclass WanVideoSamplerExtraArgs():\n @classmethod\n def INPUT_TYPES(s):\n return {\n \"required\": {\n },\n \"optional\": {\n \"riflex_freq_index\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 1000, \"step\": 1, \"tooltip\": \"Frequency index for RIFLEX, disabled when 0, default 6. Allows for new frames to be generated after without looping\"}),\n \"feta_args\": (\"FETAARGS\", ),\n \"context_options\": (\"WANVIDCONTEXT\", ),\n \"cache_args\": (\"CACHEARGS\", ),\n \"slg_args\": (\"SLGARGS\", ),\n \"rope_function\": (rope_functions, {\"default\": \"comfy\", \"tooltip\": \"Comfy's RoPE implementation doesn't use complex numbers and can thus be compiled, that should be a lot faster when using torch.compile. Chunked version has reduced peak VRAM usage when not using torch.compile\"}),\n \"loop_args\": (\"LOOPARGS\", ),\n \"experimental_args\": (\"EXPERIMENTALARGS\", ),\n \"unianimate_poses\": (\"UNIANIMATE_POSE\", ),\n \"fantasytalking_embeds\": (\"FANTASYTALKING_EMBEDS\", ),\n \"uni3c_embeds\": (\"UNI3C_EMBEDS\", ),\n \"multitalk_embeds\": (\"MULTITALK_EMBEDS\", ),\n \"guidance_mode\": (\n BERNINI_GUIDANCE_MODES,\n {\n \"default\": \"apg\",\n \"tooltip\": (\n \"Bernini guidance: none = standard CFG; apg / *_apg = Adaptive Projected \"\n \"Guidance; rv2v / v2v / v2v_chain / t2v = official modes (standard CFG in this wrapper).\"\n ),\n },\n ),\n \"apg_eta\": (\"FLOAT\", {\"default\": 0.5, \"min\": 0.0, \"max\": 1.0, \"step\": 0.01, \"tooltip\": \"APG: parallel/orthogonal balance (0=orthogonal only, 1=full)\"}),\n \"apg_momentum\": (\"FLOAT\", {\"default\": -0.5, \"min\": -1.0, \"max\": 1.0, \"step\": 0.01, \"tooltip\": \"APG: EMA momentum for smoothing guidance differences\"}),\n \"apg_norm_threshold\": (\"FLOAT\", {\"default\": 50.0, \"min\": 0.0, \"max\": 1000.0, \"step\": 1.0, \"tooltip\": \"APG: L2 norm clipping threshold (0=disabled)\"}),\n }\n }\n RETURN_TYPES = (\"WANVIDSAMPLEREXTRAARGS\",)\n RETURN_NAMES = (\"extra_args\", )\n FUNCTION = \"process\"\n CATEGORY = \"WanVideoWrapper\"\n\n def process(self, *args, **kwargs):\n return kwargs,\n\n\nclass WanVideoSamplerv2(WanVideoSampler):\n @classmethod\n def INPUT_TYPES(s):\n return {\n \"required\": {\n \"model\": (\"WANVIDEOMODEL\",),\n \"image_embeds\": (\"WANVIDIMAGE_EMBEDS\", ),\n \"cfg\": (\"FLOAT\", {\"default\": 6.0, \"min\": 0.0, \"max\": 30.0, \"step\": 0.01}),\n \"seed\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 0xffffffffffffffff}),\n \"force_offload\": (\"BOOLEAN\", {\"default\": True, \"tooltip\": \"Moves the model to the offload device after sampling\"}),\n \"scheduler\": (\"WANVIDEOSCHEDULER\",),\n },\n \"optional\": {\n \"text_embeds\": (\"WANVIDEOTEXTEMBEDS\", ),\n \"samples\": (\"LATENT\", {\"tooltip\": \"init Latents to use for video2video process\"} ),\n \"add_noise_to_samples\": (\"BOOLEAN\", {\"default\": False, \"tooltip\": \"Add noise to the samples before sampling, needed for video2video sampling when starting from clean video\"}),\n \"extra_args\": (\"WANVIDSAMPLEREXTRAARGS\", ),\n }\n }\n\n def process(self, *args, extra_args=None, **kwargs):\n import inspect\n params = inspect.signature(WanVideoSampler.process).parameters\n args_dict = {name: kwargs.get(name, param.default if param.default is not inspect.Parameter.empty else None)\n for name, param in params.items() if name != \"self\"}\n\n if extra_args is not None:\n args_dict.update(extra_args)\n else:\n args_dict[\"rope_function\"] = \"comfy\"\n\n return super().process(**args_dict)\n\n\nclass WanVideoScheduler:\n @classmethod\n def INPUT_TYPES(s):\n return {\"required\": {\n \"scheduler\": (scheduler_list, {\"default\": \"unipc\"}),\n \"steps\": (\"INT\", {\"default\": 30, \"min\": 1, \"tooltip\": \"Number of steps for the scheduler\"}),\n \"shift\": (\"FLOAT\", {\"default\": 5.0, \"min\": 0.0, \"max\": 1000.0, \"step\": 0.01}),\n \"start_step\": (\"INT\", {\"default\": 0, \"min\": 0, \"tooltip\": \"Starting step for the scheduler\"}),\n \"end_step\": (\"INT\", {\"default\": -1, \"min\": -1, \"tooltip\": \"Ending step for the scheduler\"})\n },\n \"optional\": {\n \"sigmas\": (\"SIGMAS\", ),\n \"enhance_hf\": (\"BOOLEAN\", {\"default\": False, \"tooltip\": \"Enhanced high-frequency denoising schedule\"}),\n },\n \"hidden\": {\n \"unique_id\": \"UNIQUE_ID\",\n },\n }\n\n RETURN_TYPES = (\"SIGMAS\", \"INT\", \"FLOAT\", scheduler_list, \"INT\", \"INT\",)\n RETURN_NAMES = (\"sigmas\", \"steps\", \"shift\", \"scheduler\", \"start_step\", \"end_step\")\n FUNCTION = \"process\"\n CATEGORY = \"WanVideoWrapper\"\n EXPERIMENTAL = True\n\n def process(self, scheduler, steps, start_step, end_step, shift, unique_id, sigmas=None, enhance_hf=False):\n sample_scheduler, timesteps, start_idx, end_idx = get_scheduler(\n scheduler, steps, start_step, end_step, shift, device, sigmas=sigmas, log_timesteps=True, enhance_hf=enhance_hf)\n\n scheduler_dict = {\n \"sample_scheduler\": sample_scheduler,\n \"timesteps\": timesteps,\n }\n\n _send_scheduler_sigma_plot(unique_id, sample_scheduler, start_idx, end_idx)\n\n return (sigmas, steps, shift, scheduler_dict, start_step, end_step)\n\nclass WanVideoSchedulerv2(WanVideoScheduler):\n @classmethod\n def INPUT_TYPES(s):\n import comfy.samplers\n\n return {\"required\": {\n \"sampler_name\": (\n comfy.samplers.KSampler.SAMPLERS,\n {\n \"default\": \"euler\",\n \"tooltip\": (\n \"ComfyUI KSampler sampler_name. Maps to Bernini/Wan denoise algorithm. \"\n \"Default: euler. LightX2V 4-step LoRAs: dpmpp_2m_sde.\"\n ),\n },\n ),\n \"scheduler\": (\n comfy.samplers.KSampler.SCHEDULERS,\n {\n \"default\": \"simple\",\n \"tooltip\": (\n \"ComfyUI KSampler scheduler (sigma/noise schedule). \"\n \"Default: simple. LightX2V 4-step LoRAs: sgm_uniform.\"\n ),\n },\n ),\n \"steps\": (\"INT\", {\"default\": 30, \"min\": 1, \"tooltip\": \"Number of steps for the scheduler\"}),\n \"shift\": (\"FLOAT\", {\"default\": 5.0, \"min\": 0.0, \"max\": 1000.0, \"step\": 0.01}),\n \"start_step\": (\"INT\", {\"default\": 0, \"min\": 0, \"tooltip\": \"Starting step for the scheduler\"}),\n \"end_step\": (\"INT\", {\"default\": -1, \"min\": -1, \"tooltip\": \"Ending step for the scheduler\"})\n },\n \"optional\": {\n \"sigmas\": (\"SIGMAS\", ),\n \"enhance_hf\": (\"BOOLEAN\", {\"default\": False, \"tooltip\": \"Enhanced high-frequency denoising schedule\"}),\n },\n \"hidden\": {\n \"unique_id\": \"UNIQUE_ID\",\n },\n }\n\n RETURN_TYPES = (\"WANVIDEOSCHEDULER\",)\n RETURN_NAMES = (\"scheduler\",)\n FUNCTION = \"process\"\n CATEGORY = \"WanVideoWrapper\"\n EXPERIMENTAL = True\n DESCRIPTION = (\n \"Bernini Scheduler with ComfyUI KSampler sampler_name + scheduler lists. \"\n \"Default: euler + simple. LightX2V 4-step: dpmpp_2m_sde + sgm_uniform, steps=4.\"\n )\n\n def process(\n self,\n sampler_name,\n scheduler,\n steps,\n shift,\n start_step,\n end_step,\n unique_id,\n sigmas=None,\n enhance_hf=False,\n ):\n (\n sampler_name,\n scheduler,\n steps,\n shift,\n start_step,\n end_step,\n enhance_hf,\n ) = normalize_wanvideo_schedulerv2_inputs(\n sampler_name, scheduler, steps, shift, start_step, end_step, enhance_hf\n )\n bernini_scheduler = resolve_bernini_scheduler(sampler_name)\n selected_sigmas = resolve_sigma_schedule(\n sigmas=sigmas,\n scheduler=scheduler,\n steps=int(steps),\n device=device,\n )\n sample_scheduler, timesteps, start_idx, end_idx = get_scheduler(\n bernini_scheduler,\n steps,\n start_step,\n end_step,\n shift,\n device,\n sigmas=selected_sigmas,\n log_timesteps=True,\n enhance_hf=enhance_hf,\n )\n scheduler_dict = {\n \"sample_scheduler\": sample_scheduler,\n \"timesteps\": timesteps,\n \"sigmas\": sample_scheduler.sigmas.clone(),\n \"steps\": int(steps),\n \"shift\": float(shift),\n \"start_step\": int(start_step),\n \"end_step\": int(end_step) if isinstance(end_step, int) else end_step,\n }\n _send_scheduler_sigma_plot(unique_id, sample_scheduler, start_idx, end_idx)\n return (scheduler_dict,)\n\nNODE_CLASS_MAPPINGS = {\n \"WanVideoSamplerv2\": WanVideoSamplerv2,\n \"WanVideoSamplerExtraArgs\": WanVideoSamplerExtraArgs,\n \"WanVideoSchedulerv2\": WanVideoSchedulerv2,\n }\nNODE_DISPLAY_NAME_MAPPINGS = {\n \"WanVideoSamplerv2\": \"WanVideo Sampler v2\",\n \"WanVideoSamplerExtraArgs\": \"WanVideoSampler v2 Extra Args\",\n \"WanVideoSchedulerv2\": \"WanVideo Scheduler v2\",\n }\n"}}}} |