321 lines
13 KiB
Python
321 lines
13 KiB
Python
import copy
|
|
|
|
import comfy
|
|
import torch
|
|
from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced
|
|
from comfy_extras.nodes_lt import LTXVAddGuide, LTXVCropGuides
|
|
|
|
from .latents import LTXVSelectLatents, add_latent_guide
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
@comfy_node(
|
|
name="LTXVTiledSampler",
|
|
)
|
|
class LTXVTiledSampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"vae": ("VAE",),
|
|
"noise": ("NOISE",),
|
|
"sampler": ("SAMPLER",),
|
|
"sigmas": ("SIGMAS",),
|
|
"guider": ("GUIDER",),
|
|
"latents": ("LATENT",),
|
|
"horizontal_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
|
"vertical_tiles": ("INT", {"default": 1, "min": 1, "max": 6}),
|
|
"overlap": ("INT", {"default": 1, "min": 1, "max": 8}),
|
|
"latents_cond_strength": (
|
|
"FLOAT",
|
|
{"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"boost_latent_similarity": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"crop": (["center", "disabled"], {"default": "disabled"}),
|
|
},
|
|
"optional": {
|
|
"optional_cond_images": ("IMAGE",),
|
|
"optional_cond_indices": ("STRING", {"default": "0"}),
|
|
"images_cond_strengths": ("STRING", {"default": "0.9"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"LATENT",
|
|
"LATENT",
|
|
)
|
|
RETURN_NAMES = (
|
|
"output",
|
|
"denoised_output",
|
|
)
|
|
|
|
FUNCTION = "sample"
|
|
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
vae,
|
|
noise,
|
|
sampler,
|
|
sigmas,
|
|
guider,
|
|
latents,
|
|
horizontal_tiles,
|
|
vertical_tiles,
|
|
overlap,
|
|
latents_cond_strength,
|
|
boost_latent_similarity,
|
|
crop="disabled",
|
|
optional_cond_images=None,
|
|
optional_cond_indices="0",
|
|
images_cond_strengths="0.9",
|
|
):
|
|
|
|
# Get the latent samples
|
|
samples = latents["samples"]
|
|
|
|
batch, channels, frames, height, width = samples.shape
|
|
time_scale_factor, width_scale_factor, height_scale_factor = (
|
|
vae.downscale_index_formula
|
|
)
|
|
# Validate image dimensions if provided
|
|
if optional_cond_images is not None:
|
|
img_height = height * height_scale_factor
|
|
img_width = width * width_scale_factor
|
|
cond_images = comfy.utils.common_upscale(
|
|
optional_cond_images.movedim(-1, 1),
|
|
img_width,
|
|
img_height,
|
|
"bicubic",
|
|
crop=crop,
|
|
).movedim(1, -1)
|
|
img_batch, img_height, img_width, img_channels = cond_images.shape
|
|
else:
|
|
cond_images = None
|
|
|
|
if optional_cond_indices is not None and optional_cond_images is not None:
|
|
optional_cond_indices = optional_cond_indices.split(",")
|
|
optional_cond_indices = [int(i) for i in optional_cond_indices]
|
|
assert len(optional_cond_indices) == len(
|
|
optional_cond_images
|
|
), "Number of optional cond images must match number of optional cond indices"
|
|
|
|
images_cond_strengths = [float(i) for i in images_cond_strengths.split(",")]
|
|
if optional_cond_images is not None and len(images_cond_strengths) < len(
|
|
optional_cond_images
|
|
):
|
|
# Repeat the last value to match the length of optional_cond_images
|
|
images_cond_strengths = images_cond_strengths + [
|
|
images_cond_strengths[-1]
|
|
] * (len(optional_cond_images) - len(images_cond_strengths))
|
|
|
|
# Calculate tile sizes with overlap
|
|
base_tile_height = (height + (vertical_tiles - 1) * overlap) // vertical_tiles
|
|
base_tile_width = (width + (horizontal_tiles - 1) * overlap) // horizontal_tiles
|
|
|
|
# Initialize output tensor and weight tensor
|
|
output = torch.zeros_like(samples)
|
|
denoised_output = torch.zeros_like(samples)
|
|
weights = torch.zeros_like(samples)
|
|
|
|
# Get positive and negative conditioning
|
|
try:
|
|
positive, negative = guider.raw_conds
|
|
except AttributeError:
|
|
raise ValueError(
|
|
"Guider does not have raw conds, cannot use it as a guider. "
|
|
"Please use STGGuiderAdvanced."
|
|
)
|
|
|
|
# Process each tile
|
|
for v in range(vertical_tiles):
|
|
for h in range(horizontal_tiles):
|
|
# Calculate tile boundaries
|
|
h_start = h * (base_tile_width - overlap)
|
|
v_start = v * (base_tile_height - overlap)
|
|
|
|
# Adjust end positions for edge tiles
|
|
h_end = (
|
|
min(h_start + base_tile_width, width)
|
|
if h < horizontal_tiles - 1
|
|
else width
|
|
)
|
|
v_end = (
|
|
min(v_start + base_tile_height, height)
|
|
if v < vertical_tiles - 1
|
|
else height
|
|
)
|
|
|
|
# Calculate actual tile dimensions
|
|
tile_height = v_end - v_start
|
|
tile_width = h_end - h_start
|
|
|
|
print(f"Processing tile at row {v}, col {h}:")
|
|
print(f" Position: ({v_start}:{v_end}, {h_start}:{h_end})")
|
|
print(f" Size: {tile_height}x{tile_width}")
|
|
|
|
# Extract tile
|
|
tile = samples[:, :, :, v_start:v_end, h_start:h_end]
|
|
|
|
# Create tile latents dict
|
|
tile_latents = {"samples": tile}
|
|
unconditioned_tile_latents = tile_latents.copy()
|
|
|
|
# Handle image conditioning if provided
|
|
if cond_images is not None:
|
|
# Scale coordinates for image
|
|
img_h_start = v_start * height_scale_factor
|
|
img_h_end = v_end * height_scale_factor
|
|
img_w_start = h_start * width_scale_factor
|
|
img_w_end = h_end * width_scale_factor
|
|
|
|
# Create copies of conditioning for this tile
|
|
tile_positive = positive.copy()
|
|
tile_negative = negative.copy()
|
|
|
|
for i_cond_image, (
|
|
cond_image,
|
|
cond_image_idx,
|
|
cond_image_strength,
|
|
) in enumerate(
|
|
zip(cond_images, optional_cond_indices, images_cond_strengths)
|
|
):
|
|
# Extract image tile
|
|
img_tile = cond_image[
|
|
img_h_start:img_h_end, img_w_start:img_w_end, :
|
|
].unsqueeze(0)
|
|
|
|
print(
|
|
f"Applying image conditioning on cond image {i_cond_image} for tile at row {v}, col {h} with strength {cond_image_strength} at position {cond_image_idx}:"
|
|
)
|
|
print(
|
|
f" Image tile position: ({img_h_start}:{img_h_end}, {img_w_start}:{img_w_end})"
|
|
)
|
|
print(f" Image tile size: {img_tile.shape}")
|
|
|
|
# Add guide from image tile
|
|
(
|
|
tile_positive,
|
|
tile_negative,
|
|
tile_latents,
|
|
) = LTXVAddGuide().execute(
|
|
positive=tile_positive,
|
|
negative=tile_negative,
|
|
vae=vae,
|
|
latent=tile_latents,
|
|
image=img_tile,
|
|
frame_idx=cond_image_idx,
|
|
strength=cond_image_strength,
|
|
)
|
|
if boost_latent_similarity:
|
|
middle_latent_idx = (frames - 1) // 2
|
|
middle_index_latent = LTXVSelectLatents().select_latents(
|
|
samples=unconditioned_tile_latents,
|
|
start_index=middle_latent_idx,
|
|
end_index=middle_latent_idx,
|
|
)[0]
|
|
last_index_latent = LTXVSelectLatents().select_latents(
|
|
samples=unconditioned_tile_latents,
|
|
start_index=-1,
|
|
end_index=-1,
|
|
)[0]
|
|
print(
|
|
f"using LTXVAddLatentGuide on tiled latent with latent index {middle_latent_idx} and strength {latents_cond_strength}"
|
|
)
|
|
(
|
|
tile_positive,
|
|
tile_negative,
|
|
tile_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=tile_positive,
|
|
negative=tile_negative,
|
|
latent=tile_latents,
|
|
guiding_latent=middle_index_latent,
|
|
latent_idx=middle_latent_idx,
|
|
strength=latents_cond_strength,
|
|
)
|
|
print(
|
|
f"using LTXVAddLatentGuide on tiled latent with latent index {frames-1} and strength {latents_cond_strength}"
|
|
)
|
|
(
|
|
tile_positive,
|
|
tile_negative,
|
|
tile_latents,
|
|
) = add_latent_guide(
|
|
vae=vae,
|
|
positive=tile_positive,
|
|
negative=tile_negative,
|
|
latent=tile_latents,
|
|
guiding_latent=last_index_latent,
|
|
latent_idx=frames - 1,
|
|
strength=latents_cond_strength,
|
|
)
|
|
|
|
guider = copy.copy(guider)
|
|
guider.set_conds(tile_positive, tile_negative)
|
|
|
|
# Denoise the tile
|
|
denoised_tile = SamplerCustomAdvanced().sample(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=sigmas,
|
|
latent_image=tile_latents,
|
|
)[0]
|
|
|
|
# Clean up guides if image conditioning was used
|
|
if cond_images is not None:
|
|
(
|
|
tile_positive,
|
|
tile_negative,
|
|
denoised_tile,
|
|
) = LTXVCropGuides().execute(
|
|
positive=tile_positive,
|
|
negative=tile_negative,
|
|
latent=denoised_tile,
|
|
)
|
|
|
|
# Create weight mask for this tile
|
|
tile_weights = torch.ones_like(tile)
|
|
|
|
# Apply horizontal blending weights
|
|
if h > 0: # Left overlap
|
|
h_blend = torch.linspace(0, 1, overlap, device=tile.device)
|
|
tile_weights[:, :, :, :, :overlap] *= h_blend.view(1, 1, 1, 1, -1)
|
|
if h < horizontal_tiles - 1: # Right overlap
|
|
h_blend = torch.linspace(1, 0, overlap, device=tile.device)
|
|
tile_weights[:, :, :, :, -overlap:] *= h_blend.view(1, 1, 1, 1, -1)
|
|
|
|
# Apply vertical blending weights
|
|
if v > 0: # Top overlap
|
|
v_blend = torch.linspace(0, 1, overlap, device=tile.device)
|
|
tile_weights[:, :, :, :overlap, :] *= v_blend.view(1, 1, 1, -1, 1)
|
|
if v < vertical_tiles - 1: # Bottom overlap
|
|
v_blend = torch.linspace(1, 0, overlap, device=tile.device)
|
|
tile_weights[:, :, :, -overlap:, :] *= v_blend.view(1, 1, 1, -1, 1)
|
|
|
|
# Add weighted tile to output
|
|
output[:, :, :, v_start:v_end, h_start:h_end] += (
|
|
denoised_tile["samples"] * tile_weights
|
|
)
|
|
denoised_output[:, :, :, v_start:v_end, h_start:h_end] += (
|
|
denoised_tile["samples"] * tile_weights
|
|
)
|
|
|
|
# Add weights to weight tensor
|
|
weights[:, :, :, v_start:v_end, h_start:h_end] += tile_weights
|
|
|
|
# Normalize by weights
|
|
output = output / (weights + 1e-8)
|
|
denoised_output = denoised_output / (weights + 1e-8)
|
|
|
|
return {"samples": output}, {"samples": denoised_output}
|