| import os |
| import torch |
| import numpy as np |
| from ..utils import log |
|
|
| from accelerate import init_empty_weights |
| from accelerate.utils import set_module_tensor_to_device |
|
|
| import comfy.model_management as mm |
| from comfy.utils import load_torch_file, ProgressBar |
| import folder_paths |
|
|
| script_directory = os.path.dirname(os.path.abspath(__file__)) |
| device = mm.get_torch_device() |
| offload_device = mm.unet_offload_device() |
|
|
|
|
| class WanVideoAddSteadyDancerEmbeds: |
| @classmethod |
| def INPUT_TYPES(s): |
| return {"required": { |
| "embeds": ("WANVIDIMAGE_EMBEDS",), |
| "pose_latents_positive": ("LATENT",), |
| "pose_strength_spatial": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the pose embedding"}), |
| "pose_strength_temporal": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the pose embedding"}), |
| "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the embedding application"}), |
| "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the embedding application"}), |
| }, |
| "optional": { |
| "pose_latents_negative": ("LATENT",), |
| "clip_vision_embeds": ("WANVIDIMAGE_CLIPEMBEDS",), |
| |
| } |
| } |
|
|
| RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",) |
| RETURN_NAMES = ("image_embeds",) |
| FUNCTION = "add" |
| CATEGORY = "WanVideoWrapper" |
|
|
| def add(self, embeds, pose_latents_positive, pose_strength_spatial, pose_strength_temporal, start_percent=0.0, end_percent=1.0, pose_latents_negative=None, clip_vision_embeds=None): |
| sdancer_embeds = { |
| "cond_pos": pose_latents_positive["samples"][0], |
| "cond_neg": pose_latents_negative["samples"][0] if pose_latents_negative else None, |
| "pose_strength_spatial": pose_strength_spatial, |
| "pose_strength_temporal": pose_strength_temporal, |
| "start_percent": start_percent, |
| "end_percent": end_percent, |
| "clip_fea": clip_vision_embeds, |
| } |
|
|
| updated = dict(embeds) |
| updated["sdancer_embeds"] = sdancer_embeds |
| return (updated,) |
|
|
|
|
| NODE_CLASS_MAPPINGS = { |
| "WanVideoAddSteadyDancerEmbeds": WanVideoAddSteadyDancerEmbeds, |
| } |
| NODE_DISPLAY_NAME_MAPPINGS = { |
| "WanVideoAddSteadyDancerEmbeds": "WanVideo Add SteadyDancer Embeds", |
| } |
|
|
