389 lines
14 KiB
Python
389 lines
14 KiB
Python
import torch
|
|
from torch import nn
|
|
import comfy.sample
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import gc
|
|
import logging
|
|
import nodes
|
|
|
|
class ImageMotionInfluance:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"move_range_x": ("INT", {"default": 0, "min": -150, "max": 150}),
|
|
"frame_num": ("INT", {"default": 10, "min": 2, "max": 500}),
|
|
"zoom": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.5, "step": 0.05}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "guide_motion"
|
|
CATEGORY = "image/animation"
|
|
|
|
def get_size(self, image):
|
|
image_size = image.size()
|
|
return int(image_size[2]), int(image_size[1])
|
|
|
|
def guide_motion(self, image, move_range_x, frame_num, zoom):
|
|
img_width, img_height = self.get_size(image)
|
|
|
|
step_size = abs(move_range_x) / (frame_num - 1) if move_range_x != 0 else 0
|
|
start_x = 0 if move_range_x > 0 else abs(move_range_x)
|
|
|
|
if move_range_x < 0:
|
|
start_x -= img_width
|
|
|
|
batch = []
|
|
mirrored = torch.flip(image, [2])
|
|
|
|
for i in range(frame_num):
|
|
x_pos = start_x + (step_size * i * (-1 if move_range_x < 0 else 1))
|
|
x_pos = int(x_pos)
|
|
x_pos = x_pos % img_width if move_range_x != 0 else 0
|
|
|
|
current_zoom = (i / (frame_num - 1)) * zoom if zoom > 0 else 0
|
|
if current_zoom > 0:
|
|
crop_width = int(img_width * (1 - current_zoom))
|
|
crop_height = int(img_height * (1 - current_zoom))
|
|
x_start = (img_width - crop_width) // 2
|
|
y_start = (img_height - crop_height) // 2
|
|
|
|
zoomed_original = torch.nn.functional.interpolate(
|
|
image[:, y_start:y_start + crop_height, x_start:x_start + crop_width, :].permute(0, 3, 1, 2),
|
|
size=(img_height, img_width),
|
|
mode='bilinear'
|
|
).permute(0, 2, 3, 1)
|
|
|
|
zoomed_mirror = torch.nn.functional.interpolate(
|
|
mirrored[:, y_start:y_start + crop_height, x_start:x_start + crop_width, :].permute(0, 3, 1, 2),
|
|
size=(img_height, img_width),
|
|
mode='bilinear'
|
|
).permute(0, 2, 3, 1)
|
|
else:
|
|
zoomed_original = image
|
|
zoomed_mirror = mirrored
|
|
|
|
canvas = torch.zeros((1, img_height, img_width, image.shape[3]))
|
|
|
|
remaining_width = img_width
|
|
current_x = x_pos
|
|
use_flipped = False
|
|
|
|
while remaining_width > 0:
|
|
width = min(img_width - current_x, remaining_width)
|
|
current_image = zoomed_mirror if use_flipped else zoomed_original
|
|
|
|
canvas[0, :, img_width - remaining_width:img_width - remaining_width + width, :] = \
|
|
current_image[0, :, current_x:current_x + width, :]
|
|
|
|
remaining_width -= width
|
|
current_x = 0
|
|
use_flipped = not use_flipped
|
|
|
|
batch.append(canvas)
|
|
|
|
return (torch.cat(batch, dim=0),)
|
|
|
|
class MotionGuidedSampler(nn.Module):
|
|
def __init__(
|
|
self,
|
|
motion_strength: float = 0.5,
|
|
consistency_strength: float = 0.9,
|
|
denoise_strength: float = 0.8
|
|
):
|
|
super().__init__()
|
|
self.motion_strength = motion_strength
|
|
self.consistency_strength = consistency_strength
|
|
self.denoise_strength = denoise_strength
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
def extract_motion_vector(self, current_latent: torch.Tensor, prev_latent: torch.Tensor) -> torch.Tensor:
|
|
try:
|
|
return current_latent - prev_latent
|
|
except RuntimeError as e:
|
|
self.logger.error(f"Motion extraction error: {e}")
|
|
return torch.zeros_like(current_latent)
|
|
|
|
def apply_motion_vector(self, content: torch.Tensor, motion_vector: torch.Tensor) -> torch.Tensor:
|
|
try:
|
|
motion_applied = content + (motion_vector * self.motion_strength)
|
|
if torch.isnan(motion_applied).any():
|
|
self.logger.warning("NaN detected in motion application")
|
|
return content
|
|
return motion_applied
|
|
except RuntimeError as e:
|
|
self.logger.error(f"Motion application error: {e}")
|
|
return content
|
|
|
|
def process_frames(self, latent_frames, model, positive, negative, seed, steps, cfg, sampler_name, scheduler, denoise):
|
|
batch_size = latent_frames.shape[0]
|
|
processed_frames = []
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
latent_frames = latent_frames.to(device)
|
|
pbar = comfy.utils.ProgressBar(batch_size)
|
|
|
|
# Initialize sampler
|
|
sampler = comfy.samplers.KSampler(
|
|
model,
|
|
steps=steps,
|
|
device=device,
|
|
sampler=sampler_name,
|
|
scheduler=scheduler,
|
|
denoise=denoise
|
|
)
|
|
|
|
# Process first frame
|
|
noise = comfy.sample.prepare_noise(latent_frames[0:1], seed, None).to(device)
|
|
first_frame = sampler.sample(
|
|
noise,
|
|
positive,
|
|
negative,
|
|
cfg=cfg,
|
|
latent_image=latent_frames[0:1],
|
|
force_full_denoise=True
|
|
)
|
|
|
|
processed_frames.append(first_frame)
|
|
prev_orig = latent_frames[0:1].to(device)
|
|
prev_styled = first_frame
|
|
|
|
pbar.update(1)
|
|
|
|
# Process remaining frames
|
|
for i in range(1, batch_size):
|
|
current_orig = latent_frames[i:i+1].to(device)
|
|
|
|
motion = self.extract_motion_vector(current_orig, prev_orig)
|
|
motion_guided = self.apply_motion_vector(prev_styled, motion)
|
|
|
|
current_noise = comfy.sample.prepare_noise(motion_guided, seed + i, None).to(device)
|
|
current_frame = sampler.sample(
|
|
current_noise,
|
|
positive,
|
|
negative,
|
|
cfg=cfg,
|
|
latent_image=motion_guided,
|
|
force_full_denoise=True
|
|
)
|
|
|
|
processed_frames.append(current_frame)
|
|
prev_orig = current_orig
|
|
prev_styled = current_frame
|
|
|
|
pbar.update(1)
|
|
|
|
if i % 5 == 0:
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
self.logger.info(f"Processed {i}/{batch_size} frames")
|
|
|
|
result = torch.cat(processed_frames, dim=0)
|
|
return result.to(device)
|
|
|
|
def forward(self, latent_frames, model, positive, negative, noise_seed, steps, cfg, sampler_name, scheduler, denoise):
|
|
try:
|
|
return self.process_frames(
|
|
latent_frames=latent_frames,
|
|
model=model,
|
|
positive=positive,
|
|
negative=negative,
|
|
seed=noise_seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
denoise=denoise
|
|
)
|
|
except Exception as e:
|
|
self.logger.error(f"Processing error: {str(e)}")
|
|
raise e
|
|
|
|
class HunyuanVideoSamplerSave:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"video_latents": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
|
"denoise": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(self, model, video_latents, positive, negative, seed, steps,
|
|
cfg, sampler_name, scheduler, denoise):
|
|
device = comfy.model_management.get_torch_device()
|
|
|
|
# Print latent size for debugging
|
|
print(f"Original latent size: {video_latents['samples'].size()}")
|
|
|
|
# Use nodes.common_ksampler which works with hunyuan
|
|
return nodes.common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
video_latents, # Pass the entire latent dict
|
|
denoise=denoise
|
|
)
|
|
|
|
class ResizeImageForHunyuan:
|
|
crop_methods = ["disabled", "center"]
|
|
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic"]
|
|
|
|
# Define all supported aspect ratios and their sizes
|
|
size_presets = {
|
|
"1:1": [(256,256), (384,384), (512,512), (768,768)],
|
|
"16:9": [(384,216), (512,288), (640,360), (768,432)],
|
|
"4:3": [(384,288), (512,384), (640,480), (768,576)],
|
|
"3:2": [(384,256), (512,336), (640,432), (768,512)],
|
|
"2:1": [(384,192), (512,256), (640,320), (768,384)],
|
|
"9:16": [(216,384), (288,512), (360,640), (432,768)],
|
|
"3:4": [(288,384), (384,512), (480,640), (576,768)],
|
|
"2:3": [(256,384), (336,512), (432,640), (512,768)],
|
|
"1:2": [(192,384), (256,512), (320,640), (384,768)]
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Create size options list
|
|
size_options = []
|
|
for ratio, sizes in s.size_presets.items():
|
|
for w, h in sizes:
|
|
size_options.append(f"{ratio} ({w}x{h})")
|
|
|
|
return {"required": {
|
|
"image": ("IMAGE",),
|
|
"size_preset": (size_options,),
|
|
"upscale_method": (s.upscale_methods,),
|
|
"crop": (s.crop_methods,)
|
|
}}
|
|
|
|
RETURN_TYPES = ('IMAGE',)
|
|
FUNCTION = 'resize'
|
|
CATEGORY = 'image/hunyuan'
|
|
|
|
def parse_size_preset(self, preset):
|
|
# Extract width and height from preset string like "16:9 (768x432)"
|
|
size_part = preset.split(" ")[1].strip("()")
|
|
width, height = map(int, size_part.split("x"))
|
|
return width, height
|
|
|
|
def upscale(self, image, upscale_method, width, height, crop):
|
|
samples = image.movedim(-1,1)
|
|
s = comfy.utils.common_upscale(samples, width, height, upscale_method, crop)
|
|
s = s.movedim(1,-1)
|
|
return (s,)
|
|
|
|
def resize(self, image, size_preset, upscale_method, crop):
|
|
# Get target width and height from preset
|
|
w, h = self.parse_size_preset(size_preset)
|
|
|
|
# Print info about the resize operation
|
|
print(f'Resizing image from {image.shape[2]}x{image.shape[1]} to {w}x{h} ({size_preset.split()[0]} aspect ratio)')
|
|
|
|
img = self.upscale(image, upscale_method, w, h, crop)[0]
|
|
return (img, )
|
|
|
|
class EmptyVideoLatentForHunyuan:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Define aspect ratios for Hunyuan video that can work with home GPUs (max width 768)
|
|
aspect_ratios = [
|
|
"256x256 (1:1)",
|
|
"384x384 (1:1)",
|
|
"512x512 (1:1)",
|
|
"768x768 (1:1)",
|
|
"384x216 (16:9)",
|
|
"512x288 (16:9)",
|
|
"640x360 (16:9)",
|
|
"768x432 (16:9)",
|
|
"384x288 (4:3)",
|
|
"512x384 (4:3)",
|
|
"640x480 (4:3)",
|
|
"768x576 (4:3)",
|
|
"384x256 (3:2)",
|
|
"512x336 (3:2)",
|
|
"640x432 (3:2)",
|
|
"768x512 (3:2)",
|
|
"384x192 (2:1)",
|
|
"512x256 (2:1)",
|
|
"640x320 (2:1)",
|
|
"768x384 (2:1)",
|
|
"216x384 (9:16)",
|
|
"288x512 (9:16)",
|
|
"360x640 (9:16)",
|
|
"432x768 (9:16)",
|
|
"288x384 (3:4)",
|
|
"384x512 (3:4)",
|
|
"480x640 (3:4)",
|
|
"576x768 (3:4)",
|
|
"256x384 (2:3)",
|
|
"336x512 (2:3)",
|
|
"432x640 (2:3)",
|
|
"512x768 (2:3)",
|
|
"192x384 (1:2)",
|
|
"256x512 (1:2)",
|
|
"320x640 (1:2)",
|
|
"384x768 (1:2)"
|
|
]
|
|
|
|
return {"required": {
|
|
"resolution": (aspect_ratios, ),
|
|
"length": ("INT", {"default": 25, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "latent/video"
|
|
|
|
def generate(self, resolution, length, batch_size=1):
|
|
# Parse resolution string to get dimensions
|
|
dimensions = resolution.split(' ')[0]
|
|
width, height = map(int, dimensions.split('x'))
|
|
|
|
# Ensure dimensions are multiples of 16
|
|
width = (width // 16) * 16
|
|
height = (height // 16) * 16
|
|
|
|
# Create latent with original format
|
|
latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
|
|
device=comfy.model_management.intermediate_device())
|
|
return ({"samples": latent}, )
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"HunyuanVideoSamplerSave": HunyuanVideoSamplerSave,
|
|
"ResizeImageForHunyuan": ResizeImageForHunyuan,
|
|
"EmptyVideoLatentForHunyuan": EmptyVideoLatentForHunyuan,
|
|
"ImageMotionInfluance": ImageMotionInfluance,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"HunyuanVideoSamplerSave": "Hunyuan Video Sampler Save",
|
|
"ResizeImageForHunyuan": "Resize Image For Hunyuan",
|
|
"EmptyVideoLatentForHunyuan": "Empty Video Latent For Hunyuan",
|
|
"ImageMotionInfluance": "Image Motion Influence",
|
|
}
|
|
|