| import torch
|
| from ..utils import log
|
| import comfy.model_management as mm
|
|
|
| device = mm.get_torch_device()
|
| offload_device = mm.unet_offload_device()
|
|
|
| class WanVideoAddSCAILReferenceEmbeds:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {"required": {
|
| "embeds": ("WANVIDIMAGE_EMBEDS",),
|
| "vae": ("WANVAE", {"tooltip": "VAE model"}),
|
| "ref_image": ("IMAGE",),
|
| "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the reference 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": {
|
| "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
|
| }
|
| }
|
|
|
| RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
|
| RETURN_NAMES = ("image_embeds",)
|
| FUNCTION = "add"
|
| CATEGORY = "WanVideoWrapper"
|
|
|
| def add(self, embeds, vae, ref_image, strength, start_percent, end_percent, clip_embeds=None):
|
| updated = dict(embeds)
|
|
|
| vae.to(device)
|
| ref_image_in = (ref_image[..., :3].permute(3, 0, 1, 2) * 2 - 1).to(device, vae.dtype)
|
| ref_latent = vae.encode([ref_image_in], device, tiled=False)[0]
|
| log.info(f"SCAIL ref_latent shape: {ref_latent.shape}")
|
|
|
| ref_mask = torch.ones_like(ref_latent[:4])
|
| ref_latent = torch.cat([ref_latent, ref_mask], dim=0)
|
| vae.to(offload_device)
|
|
|
| updated.setdefault("scail_embeds", {})
|
| updated["scail_embeds"]["ref_latent_pos"] = ref_latent * strength
|
| updated["scail_embeds"]["ref_latent_neg"] = torch.zeros_like(ref_latent)
|
| updated["scail_embeds"]["ref_start_percent"] = start_percent
|
| updated["scail_embeds"]["ref_end_percent"] = end_percent
|
| updated["clip_context"] = clip_embeds.get("clip_embeds", None) if clip_embeds is not None else None
|
|
|
| return (updated,)
|
|
|
| class WanVideoAddSCAILPoseEmbeds:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {"required": {
|
| "embeds": ("WANVIDIMAGE_EMBEDS",),
|
| "vae": ("WANVAE", {"tooltip": "VAE model"}),
|
| "pose_images": ("IMAGE", {"tooltip": "Pose images for the entire video"}),
|
| "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the pose control"}),
|
| "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the pose control application"}),
|
| "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the pose control application"}),
|
| },
|
| }
|
|
|
| RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
|
| RETURN_NAMES = ("image_embeds",)
|
| FUNCTION = "add"
|
| CATEGORY = "WanVideoWrapper"
|
|
|
| def add(self, embeds, vae, pose_images, strength, start_percent=0.0, end_percent=1.0):
|
| updated = dict(embeds)
|
|
|
| vae.to(device)
|
| pose_images_in = (pose_images[..., :3].permute(3, 0, 1, 2) * 2 - 1).to(device, vae.dtype)
|
| pose_latent = vae.encode([pose_images_in], device, tiled=False)[0]
|
| pose_mask = torch.ones_like(pose_latent[:4])
|
| pose_latent = torch.cat([pose_latent, pose_mask], dim=0)
|
| log.info(f"SCAIL pose_latent shape: {pose_latent.shape}")
|
|
|
| vae.to(offload_device)
|
|
|
| updated.setdefault("scail_embeds", {})
|
| updated["scail_embeds"]["pose_latent"] = pose_latent
|
| updated["scail_embeds"]["pose_strength"] = strength
|
| updated["scail_embeds"]["pose_start_percent"] = start_percent
|
| updated["scail_embeds"]["pose_end_percent"] = end_percent
|
|
|
| return (updated,)
|
|
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "WanVideoAddSCAILPoseEmbeds": WanVideoAddSCAILPoseEmbeds,
|
| "WanVideoAddSCAILReferenceEmbeds": WanVideoAddSCAILReferenceEmbeds,
|
| }
|
| NODE_DISPLAY_NAME_MAPPINGS = {
|
| "WanVideoAddSCAILReferenceEmbeds": "WanVideo Add SCAIL Reference Embeds",
|
| "WanVideoAddSCAILPoseEmbeds": "WanVideo Add SCAIL Pose Embeds",
|
| } |
