nodes/kijai-ComfyUI-DynamiCrafterWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add batch processing node

Browse Source
This commit is contained in:
kijai
2024-03-16 18:01:55 +02:00
parent 571ced1b5b
commit 610dad81b2
1 changed files with 162 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+162 -21
View File
@@ -175,25 +175,22 @@ class DynamiCrafterI2V:
#inference
ddim_sampler = DDIMSampler(self.model)
n_samples = 1
batch_variants = []
for _ in range(n_samples):
samples, _ = ddim_sampler.sample(S=steps,
conditioning=cond,
batch_size=noise_shape[0],
shape=noise_shape[1:],
verbose=False,
unconditional_guidance_scale=cfg,
unconditional_conditioning=uc,
eta=eta,
temporal_length=noise_shape[2],
conditional_guidance_scale_temporal=None,
x_T=None,
fs=fs,
timestep_spacing=timestep_spacing,
guidance_rescale=guidance_rescale,
clean_cond=True
)
samples, _ = ddim_sampler.sample(S=steps,
conditioning=cond,
batch_size=noise_shape[0],
shape=noise_shape[1:],
verbose=False,
unconditional_guidance_scale=cfg,
unconditional_conditioning=uc,
eta=eta,
temporal_length=noise_shape[2],
conditional_guidance_scale_temporal=None,
x_T=None,
fs=fs,
timestep_spacing=timestep_spacing,
guidance_rescale=guidance_rescale,
clean_cond=True
)
## reconstruct from latent to pixel space
self.model.first_stage_model.to(device)
@@ -211,14 +208,158 @@ class DynamiCrafterI2V:
last_image = video[-1].unsqueeze(0)
return (video, last_image)
class DynamiCrafterBatchInterpolation:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"model": ("DCMODEL",),
"images": ("IMAGE",),
"steps": ("INT", {"default": 50, "min": 1, "max": 200, "step": 1}),
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 20.0, "step": 0.01}),
"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.01}),
"prompt": ("STRING", {"multiline": True, "default": "",}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"fs": ("INT", {"default": 10, "min": 2, "max": 100, "step": 1}),
"dtype": (
[
'fp32',
'fp16',
], {
"default": 'fp16'
}),
"keep_model_loaded": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("IMAGE", "IMAGE",)
RETURN_NAMES = ("images", "last_image",)
FUNCTION = "process"
CATEGORY = "DynamiCrafterWrapper"
def process(self, model, images, dtype, prompt, cfg, steps, eta, seed, fs, keep_model_loaded):
device = mm.get_torch_device()
mm.unload_all_models()
mm.soft_empty_cache()
torch.manual_seed(seed)
dtype = model.dtype
self.model = model
channels = self.model.model.diffusion_model.out_channels
frames = self.model.temporal_length
images = images * 2 - 1
images = images.permute(0, 3, 1, 2).to(dtype).to(device)
out = []
autocast_condition = (dtype != torch.float32) and not comfy.model_management.is_device_mps(device)
with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext():
for i in range(len(images) - 1):
image = images[i].unsqueeze(0)
image2 = images[i+1].unsqueeze(0)
B, C, H, W = image.shape
noise_shape = [B, channels, frames, H // 8, W // 8]
self.model.first_stage_model.to(device)
z = get_latent_z(self.model, image.unsqueeze(2)) #bc,1,hw
z2 = get_latent_z(self.model, image2.unsqueeze(2)) #bc,1,hw
img_tensor_repeat = repeat(z, 'b c t h w -> b c (repeat t) h w', repeat=frames)
img_tensor_repeat = torch.zeros_like(img_tensor_repeat)
img_tensor_repeat[:,:,:1,:,:] = z
img_tensor_repeat[:,:,-1:,:,:] = z2
self.model.first_stage_model.to('cpu')
self.model.cond_stage_model.to(device)
self.model.embedder.to(device)
self.model.image_proj_model.to(device)
text_emb = self.model.get_learned_conditioning([prompt])
cond_images = self.model.embedder(image)
img_emb = self.model.image_proj_model(cond_images)
imtext_cond = torch.cat([text_emb, img_emb], dim=1)
fs = torch.tensor([fs], dtype=torch.long, device=self.model.device)
cond = {"c_crossattn": [imtext_cond], "c_concat": [img_tensor_repeat]}
if noise_shape[-1] == 32:
timestep_spacing = "uniform"
guidance_rescale = 0.0
else:
timestep_spacing = "uniform_trailing"
guidance_rescale = 0.7
## construct unconditional guidance
if cfg != 1.0:
uc_emb = self.model.get_learned_conditioning([""])
## process image embedding token
if hasattr(self.model, 'embedder'):
uc_img = torch.zeros(noise_shape[0],3,224,224).to(self.model.device)
## img: b c h w >> b l c
uc_img = self.model.embedder(uc_img)
uc_img = self.model.image_proj_model(uc_img)
uc_emb = torch.cat([uc_emb, uc_img], dim=1)
if isinstance(cond, dict):
uc = {key:cond[key] for key in cond.keys()}
uc.update({'c_crossattn': [uc_emb]})
else:
uc = uc_emb
else:
uc = None
self.model.cond_stage_model.to('cpu')
self.model.embedder.to('cpu')
self.model.image_proj_model.to('cpu')
#inference
ddim_sampler = DDIMSampler(self.model)
samples, _ = ddim_sampler.sample(S=steps,
conditioning=cond,
batch_size=noise_shape[0],
shape=noise_shape[1:],
verbose=False,
unconditional_guidance_scale=cfg,
unconditional_conditioning=uc,
eta=eta,
temporal_length=noise_shape[2],
conditional_guidance_scale_temporal=None,
x_T=None,
fs=fs,
timestep_spacing=timestep_spacing,
guidance_rescale=guidance_rescale,
clean_cond=True
)
## reconstruct from latent to pixel space
self.model.first_stage_model.to(device)
decoded_images = self.model.decode_first_stage(samples) #b c t h w
self.model.first_stage_model.to('cpu')
video = decoded_images.detach().cpu()
video = torch.clamp(video.float(), -1., 1.)
video = (video + 1.0) / 2.0
video = video.squeeze(0).permute(1, 2, 3, 0)
print(f"Sampled {i} / {len(images) - 1}")
out.append(video)
if not keep_model_loaded:
self.model = None
mm.soft_empty_cache()
out_video = torch.cat(out, dim=0)
last_image = out_video[-1].unsqueeze(0)
return (out_video, last_image)
NODE_CLASS_MAPPINGS = {
"DynamiCrafterI2V": DynamiCrafterI2V,
"DynamiCrafterModelLoader": DynamiCrafterModelLoader
"DynamiCrafterModelLoader": DynamiCrafterModelLoader,
"DynamiCrafterBatchInterpolation": DynamiCrafterBatchInterpolation
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DynamiCrafterI2V": "DynamiCrafterI2V",
"DynamiCrafterModelLoader": "DynamiCrafterModelLoader"
"DynamiCrafterModelLoader": "DynamiCrafterModelLoader",
"DynamiCrafterBatchInterpolation": "DynamiCrafterBatchInterpolation"
}