147 lines
5.5 KiB
Python
147 lines
5.5 KiB
Python
"""MiniMax H3 Reference to Video (Lam) - 图片数组/音频数组/音频数组 input variant."""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torchaudio
|
|
|
|
import comfy.model_management
|
|
import comfy.nested_tensor
|
|
import comfy.utils
|
|
import node_helpers
|
|
from comfy.comfy_types import ComfyNodeABC
|
|
|
|
CANVAS_MULTIPLE = 32
|
|
BASE_SHORT_EDGE = 768
|
|
MAX_PIXELS = 768 * 1344
|
|
REF_IMAGE_SHORT_EDGE = 2048
|
|
FPS = 24
|
|
AUDIO_LATENT_FPS = 40
|
|
|
|
|
|
def align_frame_count(n):
|
|
while n % 17 != 5:
|
|
n += 1
|
|
return n
|
|
|
|
|
|
def video_latent_t(frame_count):
|
|
return 2 if frame_count <= 5 else ((frame_count - 5) // 17) * 5 + 2
|
|
|
|
|
|
def temporal_shape(length):
|
|
frame_count = align_frame_count(max(5, length))
|
|
duration = frame_count / FPS
|
|
return frame_count, video_latent_t(frame_count), round(duration * AUDIO_LATENT_FPS)
|
|
|
|
|
|
def _resize(image, width, height, crop):
|
|
samples = image[..., :3].movedim(-1, 1)
|
|
samples = comfy.utils.common_upscale(samples, width, height, "lanczos", crop)
|
|
return samples.movedim(1, -1)
|
|
|
|
|
|
def _empty_av_latent(width, height, length, batch_size=1):
|
|
frame_count, latent_t, audio_t = temporal_shape(length)
|
|
video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
|
|
device=comfy.model_management.intermediate_device())
|
|
audio = torch.zeros([batch_size, 32, 2, audio_t],
|
|
device=comfy.model_management.intermediate_device())
|
|
return {"samples": comfy.nested_tensor.NestedTensor((video, audio))}, frame_count
|
|
|
|
|
|
class LamMiniMaxH3RefVideo(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP", {}),
|
|
"vae": ("VAE", {}),
|
|
"audio_vae": ("VAE", {}),
|
|
"prompt": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
"width": ("INT", {"default": 1344, "min": 32, "max": 8192, "step": 32}),
|
|
"height": ("INT", {"default": 768, "min": 32, "max": 8192, "step": 32}),
|
|
"length": ("INT", {"default": 124, "min": 5, "max": 3600, "step": 17}),
|
|
"ref_image_size": (["match", "max"], {"default": "match"}),
|
|
},
|
|
"optional": {
|
|
"ref_images": ("IMAGE", {}),
|
|
"ref_video_audios": ("AUDIO", {}),
|
|
"ref_audios": ("AUDIO", {}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "latent")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "lam"
|
|
DESCRIPTION = "MiniMax H3 图片+音频 参视频节点"
|
|
|
|
@staticmethod
|
|
def _encode_ref_audio(audio_vae, audio):
|
|
waveform = audio["waveform"] # [B, C, L]
|
|
sr = audio["sample_rate"]
|
|
vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
|
|
if sr != vae_sr:
|
|
waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
|
|
z = audio_vae.encode(waveform[:1].movedim(1, -1)) # [1, 32, 2, T]
|
|
return z, z.shape[-1]
|
|
|
|
def execute(self, clip, vae, audio_vae, prompt, width, height, length,
|
|
ref_image_size="match", ref_images=None,
|
|
ref_video_audios=None, ref_audios=None):
|
|
latent, frame_count = _empty_av_latent(width, height, length)
|
|
|
|
ref_items = []
|
|
ref_blocks = []
|
|
|
|
# 图片数组
|
|
if ref_images is not None:
|
|
if ref_image_size == "match":
|
|
scale = min(1.0, math.sqrt((width * height) / (ref_images.shape[1] * ref_images.shape[2])))
|
|
else:
|
|
scale = min(1.0, REF_IMAGE_SHORT_EDGE / min(ref_images.shape[1], ref_images.shape[2]))
|
|
tw = max(CANVAS_MULTIPLE, round(ref_images.shape[2] * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
|
|
th = max(CANVAS_MULTIPLE, round(ref_images.shape[1] * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
|
|
resized = _resize(ref_images, tw, th, "disabled")
|
|
z = vae.encode(resized)
|
|
ref_items.append({"type": "image", "data": resized})
|
|
ref_blocks.append({"kind": "image", "latent_h": th // 16, "latent_w": tw // 16, "latent": z})
|
|
|
|
# 音频数组 1: 视频伴音 (ref_video_audios)
|
|
if ref_video_audios is not None:
|
|
if not isinstance(ref_video_audios, (list, tuple)):
|
|
ref_video_audios = [ref_video_audios]
|
|
for audio in ref_video_audios:
|
|
if audio is None:
|
|
continue
|
|
audio_latent, ref_audio_t = self._encode_ref_audio(audio_vae, audio)
|
|
ref_items.append({"type": "audio"})
|
|
ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
|
|
|
|
# 音频数组 2: 独立音频 (ref_audios)
|
|
if ref_audios is not None:
|
|
if not isinstance(ref_audios, (list, tuple)):
|
|
ref_audios = [ref_audios]
|
|
for audio in ref_audios:
|
|
if audio is None:
|
|
continue
|
|
audio_latent, ref_audio_t = self._encode_ref_audio(audio_vae, audio)
|
|
ref_items.append({"type": "audio"})
|
|
ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
|
|
|
|
tokens = clip.tokenize(prompt, minimax_ref_items=ref_items)
|
|
cond = clip.encode_from_tokens_scheduled(tokens)
|
|
if ref_blocks:
|
|
cond = node_helpers.conditioning_set_values(cond, {"minimax_refs": ref_blocks})
|
|
|
|
return (cond, latent)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LamMiniMaxH3RefVideo": LamMiniMaxH3RefVideo,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LamMiniMaxH3RefVideo": "Lam MiniMax H3 Reference to Video",
|
|
}
|