96 lines
4.3 KiB
Python
96 lines
4.3 KiB
Python
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",
|
|
} |