import torch import comfy.sample import comfy.model_management from comfy.utils import ProgressBar from .utils import parse_tile_prompts from .utils import apply_controlnet_to_conditioning, blend_and_place_tile, build_working_tensor def _apply_zimage_patch(model, model_patch, vae, working_tensor, keep_mask, strength): """Patch the model with the DiffSynth/Fun inpaint control (z-image, Qwen) for one tile. keep_mask: [H, W] with 1 where neighbour pixels are kept, 0 where new content is generated. The node's convention is white(1) = regenerate, so we pass its inverse. """ from comfy_extras.nodes_model_patch import ZImageFunControlnet gen_mask = (1.0 - keep_mask).unsqueeze(0) # (1, H, W); white = generate return ZImageFunControlnet().diffsynth_controlnet( model=model, model_patch=model_patch, vae=vae, inpaint_image=working_tensor, strength=strength, mask=gen_mask, )[0] def _decode_tile(vae, samples): img = vae.decode(samples) if img.ndim == 5: # video VAE (Krea2/Wan) -> single frame img = img[:, 0] return img[0].cpu() # [gen_h8, gen_w8, 3] class TiledImageGenerator: """ LLM-directed tiled generation. Each tile generates on an expanded canvas that reaches into its already-generated neighbours' overlap; those neighbour pixels drive coherence via whichever mechanism is wired: * controlnet -> ControlNet on the conditioning (SDXL, Flux/Krea2) * model_patch -> DiffSynth/Fun inpaint patch on the model (z-image, Qwen) With neither wired, tiles generate independently (no coherence). Results are feather-placed into one canvas. """ @classmethod def INPUT_TYPES(cls): import comfy.samplers return { "required": { "json_tile_prompts": ("STRING", {"multiline": True}), "global_positive": ("STRING", {"multiline": True, "default": ""}), "global_negative": ("STRING", {"multiline": True, "default": ""}), "grid_width": ("INT", {"default": 4, "min": 1, "max": 16}), "grid_height": ("INT", {"default": 6, "min": 1, "max": 16}), "tile_width": ("INT", {"default": 1024, "min": 256, "max": 2048}), "tile_height": ("INT", {"default": 1024, "min": 256, "max": 2048}), "overlap_percent": ("FLOAT", {"default": 0.15, "min": 0.05, "max": 0.5, "step": 0.01}), "model": ("MODEL",), "clip": ("CLIP",), "vae": ("VAE",), "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), "steps": ("INT", {"default": 20, "min": 1, "max": 100}), "cfg": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1}), "control_strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the coherence control (ControlNet or model patch)."}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}), "seamlessX": ("BOOLEAN", {"default": True, "tooltip": "If true, side of image will be seamless. (2+ tiles)"}), "seamlessY": ("BOOLEAN", {"default": False, "tooltip": "If true, top/bottom of image will be seamless. (2+ tiles)"}), }, "optional": { "controlnet": ("CONTROL_NET",), "model_patch": ("MODEL_PATCH",), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("composite_image", "individual_tiles") FUNCTION = "generate_tiled_image" CATEGORY = "image/generation" def generate_tiled_image(self, json_tile_prompts, global_positive, global_negative, grid_width, grid_height, tile_width, tile_height, overlap_percent, seed, model, clip, vae, sampler_name, scheduler, steps, cfg, control_strength, seamlessX, seamlessY, controlnet=None, model_patch=None, control_net=None): # Accept the older `control_net` input name from previously-wired graphs. if controlnet is None: controlnet = control_net tile_prompts = parse_tile_prompts(json_tile_prompts, grid_width, grid_height) overlap_x = int(tile_width * overlap_percent) overlap_y = int(tile_height * overlap_percent) final_width = tile_width * grid_width final_height = tile_height * grid_height mode = ("model_patch" if model_patch is not None else "controlnet" if controlnet is not None else "none") print(f"Final output size: {final_width} x {final_height}") print(f"Tile size: {tile_width} x {tile_height} | Grid: {grid_width} x {grid_height}") print(f"Overlap: {overlap_x} x {overlap_y} pixels ({overlap_percent * 100}%) | coherence: {mode}") final_tensor = torch.zeros((1, final_height, final_width, 3), dtype=torch.float32) individual_tiles = [] device = model.load_device pbar = ProgressBar(grid_width * grid_height) for y in range(grid_height): for x in range(grid_width): idx = y * grid_width + x current_prompt = tile_prompts[idx] current_seed = seed + idx final_pos_x = x * tile_width final_pos_y = y * tile_height has_left_neighbor = x > 0 has_top_neighbor = y > 0 # Both coherence mechanisms need the neighbour pixels, so both use the # expanded canvas + working tensor. coherence_active = controlnet is not None or model_patch is not None if coherence_active: gen_w = (tile_width + (overlap_x if has_left_neighbor else 0) + (overlap_x if seamlessX and x == grid_width - 1 else 0)) gen_h = (tile_height + (overlap_y if has_top_neighbor else 0) + (overlap_y if seamlessY and y == grid_height - 1 else 0)) else: gen_w, gen_h = tile_width, tile_height gen_w8 = ((gen_w + 7) // 8) * 8 gen_h8 = ((gen_h + 7) // 8) * 8 print(f"Tile ({x+1},{y+1}) gen canvas: {gen_w8}x{gen_h8}, canvas pos: ({final_pos_x},{final_pos_y})") print(f"Generating tile ({x+1},{y+1}) prompt: {current_prompt}") # Per-tile encoding: a shared neg_cond corrupts SDXL pooled_output across tiles. combined_prompt = f"{current_prompt} {global_positive}" if global_positive else current_prompt pos_cond = clip.encode_from_tokens_scheduled(clip.tokenize(combined_prompt)) neg_cond = clip.encode_from_tokens_scheduled(clip.tokenize(global_negative)) positive, negative = pos_cond, neg_cond tile_model = model working_tensor = None keep_mask = None if coherence_active: working_tensor, keep_mask = build_working_tensor( final_tensor, final_pos_x, final_pos_y, tile_width, tile_height, overlap_x, overlap_y, gen_w8, gen_h8, has_left_neighbor, has_top_neighbor, wrap_x=(seamlessX and x == grid_width - 1), wrap_y=(seamlessY and y == grid_height - 1)) if model_patch is not None: tile_model = _apply_zimage_patch( model, model_patch, vae, working_tensor, keep_mask, control_strength) elif controlnet is not None: positive, negative = apply_controlnet_to_conditioning( positive=pos_cond, negative=neg_cond, control_net=controlnet, image=working_tensor, strength=control_strength, start_percent=0.0, end_percent=1.0, vae=vae) # 16-channel (z-image/Krea2) and video-latent formats need the right shape. latent_image = comfy.sample.fix_empty_latent_channels( model, torch.zeros((1, 4, gen_h8 // 8, gen_w8 // 8), device=device)) noise = comfy.sample.prepare_noise(latent_image, current_seed, None) samples = comfy.sample.sample( tile_model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image) decoded = _decode_tile(vae, samples) # [gen_h8, gen_w8, 3] # For the model_patch path the control's reconstruction of the kept # overlap can return flat/grey; force those pixels back to the real # neighbour so the feather blends real content, not grey. if model_patch is not None and keep_mask is not None: k = keep_mask.unsqueeze(-1) decoded = (1.0 - k) * decoded + k * working_tensor[0] start_x = (overlap_x if has_left_neighbor else 0) if coherence_active else 0 start_y = (overlap_y if has_top_neighbor else 0) if coherence_active else 0 individual_tiles.append( decoded[start_y:start_y + tile_height, start_x:start_x + tile_width, :].unsqueeze(0)) blend_and_place_tile( final_tensor, decoded, final_pos_x, final_pos_y, tile_width, tile_height, overlap_x, overlap_y, has_left_neighbor, has_top_neighbor, coherence_active) # Flushing the CUDA cache every tile costs real time and mostly frees # memory the next tile immediately re-allocates; throttle it. if (idx + 1) % 4 == 0: comfy.model_management.soft_empty_cache() pbar.update(1) comfy.model_management.soft_empty_cache() tile_batch = torch.cat(individual_tiles, dim=0) if individual_tiles else \ torch.zeros((1, tile_height, tile_width, 3), dtype=torch.float32) if seamlessX and overlap_x > 0: final_tensor = final_tensor[:, :, :final_width - overlap_x, :] if seamlessY and overlap_y > 0: final_tensor = final_tensor[:, :final_height - overlap_y, :, :] return final_tensor, tile_batch NODE_CLASS_MAPPINGS = { "TiledImageGenerator": TiledImageGenerator, } NODE_DISPLAY_NAME_MAPPINGS = { "TiledImageGenerator": "Tiled Image Generator", }