vae tiled fixes and few other mistakes

comfy/model_management.py

@@ -800,7 +800,7 @@ class LoadedModel:
             return self.model_offloaded_memory()
 
         # Handle AutoencoderKL
-        if self.model.model is not None and isinstance(self.model.model, AutoencoderKL):
+        if self.model is not None and isinstance(self.model.model, AutoencoderKL):
             shape = getattr(self.model, 'last_shape', (1, 4, 64, 64))
             dtype = getattr(self.model, 'model_dtype', torch.float32)()
             return estimate_vae_decode_memory(self.model.model, shape, dtype)

comfy/sd.py

@@ -28,6 +28,7 @@ from . import model_detection
 
 from . import sd1_clip
 from . import sdxl_clip
+from comfy.cli_args import args
 import comfy.text_encoders.sd2_clip
 import comfy.text_encoders.sd3_clip
 import comfy.text_encoders.sa_t5
@@ -54,6 +55,8 @@ import comfy.taesd.taesd
 
 import comfy.ldm.flux.redux
 
+DEBUG_ENABLED = args.debug
+
 def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
     key_map = {}
     if model is not None:
@@ -497,18 +500,23 @@ class VAE:
                 pixels = pixels.narrow(d + 1, x_offset, x)
         return pixels
 
-    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
+    def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap=16):
+        # Calculate progress bar steps for a single pass
         steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
-        steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
-        steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
         pbar = comfy.utils.ProgressBar(steps)
 
-        decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
-        output = self.process_output(
-            (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-            comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-             comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar))
-            / 3.0)
+        # Define decode function with tile logging
+        if not DEBUG_ENABLED:
+            decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
+        else:
+            decode_fn = lambda a: (logging.debug(f"Tile shape: {a.shape}, min: {a.min()}, max: {a.max()}"),
+                                   self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float())[1]
+
+        # Single pass with provided tile sizes
+        output = comfy.utils.tiled_scale(
+            samples, decode_fn, tile_x, tile_y, overlap,
+            upscale_amount=self.upscale_ratio, output_device=self.output_device, pbar=pbar
+        )
         return output
 
     def decode_tiled_1d(self, samples, tile_x=128, overlap=32):
@@ -1068,9 +1076,10 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
             else:
                 logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
 
-    left_over = sd.keys()
-    if len(left_over) > 0:
-        logging.debug("left over keys: {}".format(left_over))
+    if DEBUG_ENABLED:
+        left_over = sd.keys()
+        if len(left_over) > 0:
+                logging.debug("left over keys: {}".format(left_over))
 
     if output_model:
         model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
@@ -1137,9 +1146,10 @@ def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffuse
     model = model_config.get_model(new_sd, "")
     model = model.to(offload_device)
     model.load_model_weights(new_sd, "")
-    left_over = sd.keys()
-    if len(left_over) > 0:
-        logging.info("left over keys in unet: {}".format(left_over))
+    if DEBUG_ENABLED:
+        left_over = sd.keys()
+        if len(left_over) > 0:
+            logging.info("left over keys in unet: {}".format(left_over))
     return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
 
 

comfy/utils.py

@@ -33,6 +33,9 @@ from comfy.cli_args import args
 MMAP_TORCH_FILES = args.mmap_torch_files
 
 ALWAYS_SAFE_LOAD = False
+
+DEBUG_ENABLED = args.debug
+
 if hasattr(torch.serialization, "add_safe_globals"):  # TODO: this was added in pytorch 2.4, the unsafe path should be removed once earlier versions are deprecated
     class ModelCheckpoint:
         pass
@@ -876,116 +879,205 @@ def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):
 
 @torch.inference_mode()
 def tiled_scale_multidim(samples, function, tile=(64, 64), overlap=8, upscale_amount=4, out_channels=3, output_device="cpu", downscale=False, index_formulas=None, pbar=None):
+    """
+    Perform tiled scaling of input samples using the provided function with overlap blending.
+    
+    Args:
+        samples: Input tensor of shape [batch, channels, *spatial_dims].
+        function: Function to process each tile (e.g., VAE decode).
+        tile: Tuple of tile sizes for each spatial dimension.
+        overlap: Overlap size or list of overlaps for each dimension.
+        upscale_amount: Scaling factor or list of factors for each dimension.
+        out_channels: Number of output channels.
+        output_device: Device for output tensor.
+        downscale: If True, downscale instead of upscale.
+        index_formulas: Scaling factors for tile positions (defaults to upscale_amount).
+        pbar: Optional progress bar object.
+    
+    Returns:
+        Scaled output tensor of shape [batch, out_channels, *scaled_spatial_dims].
+    """
     dims = len(tile)
 
-    if not (isinstance(upscale_amount, (tuple, list))):
+    # Wrap function to ensure FP32
+    def fp32_function(x):
+        x_fp32 = x.to(dtype=torch.float32)
+        result = function(x_fp32)
+        return result.to(dtype=torch.float32)
+
+    # Convert parameters to lists for multidimensional support
+    if not isinstance(upscale_amount, (tuple, list)):
         upscale_amount = [upscale_amount] * dims
-
-    if not (isinstance(overlap, (tuple, list))):
+    if not isinstance(overlap, (tuple, list)):
         overlap = [overlap] * dims
-
     if index_formulas is None:
         index_formulas = upscale_amount
-
-    if not (isinstance(index_formulas, (tuple, list))):
+    if not isinstance(index_formulas, (tuple, list)):
         index_formulas = [index_formulas] * dims
 
+    # Define scaling functions
     def get_upscale(dim, val):
         up = upscale_amount[dim]
-        if callable(up):
-            return up(val)
-        else:
-            return up * val
+        return up(val) if callable(up) else up * val
 
     def get_downscale(dim, val):
         up = upscale_amount[dim]
-        if callable(up):
-            return up(val)
-        else:
-            return val / up
+        return up(val) if callable(up) else val / up
 
     def get_upscale_pos(dim, val):
         up = index_formulas[dim]
-        if callable(up):
-            return up(val)
-        else:
-            return up * val
+        return up(val) if callable(up) else up * val
 
     def get_downscale_pos(dim, val):
         up = index_formulas[dim]
-        if callable(up):
-            return up(val)
-        else:
-            return val / up
+        return up(val) if callable(up) else val / up
 
-    if downscale:
-        get_scale = get_downscale
-        get_pos = get_downscale_pos
-    else:
-        get_scale = get_upscale
-        get_pos = get_upscale_pos
+    get_scale = get_downscale if downscale else get_upscale
+    get_pos = get_downscale_pos if downscale else get_upscale_pos
 
     def mult_list_upscale(a):
-        out = []
-        for i in range(len(a)):
-            out.append(round(get_scale(i, a[i])))
-        return out
+        """Compute scaled dimensions for output tensor."""
+        return [round(get_scale(i, a[i])) for i in range(len(a))]
 
-    output = torch.empty([samples.shape[0], out_channels] + mult_list_upscale(samples.shape[2:]), device=output_device)
+    # Initialize output tensor
+    output_shape = [samples.shape[0], out_channels] + mult_list_upscale(samples.shape[2:])
+    output = torch.empty(output_shape, device=output_device, dtype=torch.float32)
+    if DEBUG_ENABLED:
+        logging.debug(f"Input shape: {samples.shape}, output shape: {output_shape}, tile: {tile}")
+        logging.debug(f"Input stats: min={samples.min():.4f}, max={samples.max():.4f}, mean={samples.mean():.4f}")
+
+    # Test VAE
+    try:
+        test_input = samples[:1].to(dtype=torch.float32)
+        test_result = fp32_function(test_input).to(output_device, dtype=torch.float32)
+        if DEBUG_ENABLED:
+            logging.debug(f"VAE test result: shape={test_result.shape}, min={test_result.min():.4f}, max={test_result.max():.4f}")
+        if torch.isnan(test_result).any() or torch.isinf(test_result).any():
+            logging.error("VAE produces NaN or Inf in test output. Check VAE model or input latents.")
+            raise RuntimeError("VAE output contains NaN or Inf")
+    except Exception as e:
+        logging.error(f"VAE test failed: {e}")
+        raise
 
     for b in range(samples.shape[0]):
         s = samples[b:b+1]
 
-        # handle entire input fitting in a single tile
+        # Handle case where input fits in a single tile
         if all(s.shape[d+2] <= tile[d] for d in range(dims)):
-            output[b:b+1] = function(s).to(output_device)
+            s_fp32 = s.to(dtype=torch.float32)
+            result = fp32_function(s_fp32).to(output_device, dtype=torch.float32)
+            if DEBUG_ENABLED:
+                logging.debug(f"Single tile result: shape={result.shape}, min={result.min():.4f}, max={result.max():.4f}")
+            if result.shape == output_shape[1:]:
+                output[b:b+1] = result
+            else:
+                result = result.narrow(1, 0, output_shape[1])
+                for d in range(dims):
+                    result = result.narrow(d + 2, 0, output_shape[d + 2])
+                output[b:b+1] = result
             if pbar is not None:
                 pbar.update(1)
             continue
 
-        out = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)
-        out_div = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)
+        # Initialize accumulation tensors
+        out = torch.zeros(output_shape[1:], device=output_device, dtype=torch.float32)
+        out_div = torch.full_like(out, 1e-6)
 
-        positions = [range(0, s.shape[d+2] - overlap[d], tile[d] - overlap[d]) if s.shape[d+2] > tile[d] else [0] for d in range(dims)]
+        # Compute tile positions
+        positions = []
+        tile_counts = []
+        for d in range(dims):
+            step = max(1, tile[d] - overlap[d])
+            end = max(0, s.shape[d+2] - tile[d])
+            pos = list(range(0, end + 1, step))
+            if pos and (pos[-1] < end or s.shape[d+2] > tile[d]):
+                pos.append(end)
+            positions.append(pos if pos else [0])
+            tile_counts.append(len(pos))
+        if DEBUG_ENABLED:
+            logging.debug(f"Tile positions: {positions}")
 
+        # Process each tile
+        total_tiles = max(1, len(list(itertools.product(*positions))))
         for it in itertools.product(*positions):
             s_in = s
             upscaled = []
 
+            # Extract tile
             for d in range(dims):
-                pos = max(0, min(s.shape[d + 2] - overlap[d], it[d]))
-                l = min(tile[d], s.shape[d + 2] - pos)
-                s_in = s_in.narrow(d + 2, pos, l)
+                pos = max(0, min(s.shape[d+2] - tile[d], it[d]))
+                length = min(tile[d], s.shape[d+2] - pos)
+                s_in = s_in.narrow(d + 2, pos, length)
                 upscaled.append(round(get_pos(d, pos)))
 
-            ps = function(s_in).to(output_device)
-            mask = torch.ones_like(ps)
+            # Process tile
+            s_in_fp32 = s_in.to(dtype=torch.float32)
+            ps = fp32_function(s_in_fp32).to(output_device, dtype=torch.float32)
+            if DEBUG_ENABLED:
+                logging.debug(f"Tile at {it}: input={s_in.shape}, output={ps.shape}, min={ps.min():.4f}, max={ps.max():.4f}")
+            if torch.isnan(ps).any() or torch.isinf(ps).any():
+                if DEBUG_ENABLED:
+                    logging.warning(f"Tile at {it} contains NaN or Inf, clamping values")
+                ps = torch.clamp(ps, min=-1e6, max=1e6)
+            mask = torch.ones_like(ps, dtype=torch.float32) / total_tiles
 
-            for d in range(2, dims + 2):
-                feather = round(get_scale(d - 2, overlap[d - 2]))
-                if feather >= mask.shape[d]:
+            # Apply feathering for smooth overlap blending
+            for d in range(dims):
+                feather = min(round(get_scale(d, overlap[d])), ps.shape[d + 2] // 16)
+                if feather < 1:
                     continue
                 for t in range(feather):
                     a = (t + 1) / feather
-                    mask.narrow(d, t, 1).mul_(a)
-                    mask.narrow(d, mask.shape[d] - 1 - t, 1).mul_(a)
+                    mask.narrow(d + 2, t, 1).mul_(a)
+                    mask.narrow(d + 2, mask.shape[d + 2] - 1 - t, 1).mul_(a)
+            mask = mask.clamp(min=1e-6)
 
+            # Accumulate results
             o = out
             o_d = out_div
             for d in range(dims):
-                o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2])
-                o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2])
+                start = upscaled[d]
+                size = min(ps.shape[d + 2], output_shape[d + 2] - start)
+                o = o.narrow(d + 1, start, size)
+                o_d = o_d.narrow(d + 1, start, size)
+                ps = ps.narrow(d + 2, 0, size)
+                mask = mask.narrow(d + 2, 0, size)
 
+            # Squeeze batch dimension
+            ps = ps.squeeze(0)
+            mask = mask.squeeze(0)
             o.add_(ps * mask)
             o_d.add_(mask)
 
             if pbar is not None:
                 pbar.update(1)
 
-        output[b:b+1] = out/out_div
-    return output
+        # Fallback to non-tiled if NaN
+        if torch.isnan(out).any():
+            if DEBUG_ENABLED:
+                logging.warning("NaN detected in tiled output, falling back to non-tiled")
+            s_fp32 = s.to(dtype=torch.float32)
+            result = fp32_function(s_fp32).to(output_device, dtype=torch.float32)
+            if result.shape == output_shape[1:]:
+                output[b:b+1] = result
+            else:
+                result = result.narrow(1, 0, output_shape[1])
+                for d in range(dims):
+                    result = result.narrow(d + 2, 0, output_shape[d + 2])
+                output[b:b+1] = result
+        else:
+            if DEBUG_ENABLED:
+                logging.debug(f"out stats: min={out.min():.4f}, max={out.max():.4f}")
+                logging.debug(f"out_div stats: min={out_div.min():.4f}, max={out_div.max():.4f}")
+            output[b:b+1] = out / out_div
 
-def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
+        if pbar is not None:
+            pbar.update(1)
+
+    return output
+    
+def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap=8, upscale_amount=4, out_channels=3, output_device="cpu", pbar=None):
+    """Wrapper for 2D tiled scaling."""
     return tiled_scale_multidim(samples, function, (tile_y, tile_x), overlap=overlap, upscale_amount=upscale_amount, out_channels=out_channels, output_device=output_device, pbar=pbar)
 
 PROGRESS_BAR_ENABLED = True

fast_sampler.py

@@ -8,6 +8,7 @@ from comfy.model_management import get_torch_device, vae_dtype, soft_empty_cache
 from contextlib import contextmanager
 import latent_preview
 import logging
+import traceback
 
 # Global flag for profiling
 PROFILING_ENABLED = args.profile
@@ -41,15 +42,18 @@ def profile_cuda_sync(is_gpu, message="CUDA sync"):
         logging.debug(f"{message} took {time.time() - sync_start:.3f} s")
 
 def is_fp16_safe(device):
-    """Check if FP16 is safe for the GPU (disabled for GTX 1660/Turing)."""
+    """Check if FP16 is safe for the GPU (disabled for Turing)."""
     if device.type != 'cuda':
         return False
     if device in _fp16_safe_cache:
         return _fp16_safe_cache[device]
     try:
         props = torch.cuda.get_device_properties(device)
-        is_safe = props.major >= 8 or props.compute_capability[0] > 7
+        # Disable FP16 for Turing (major == 7) and earlier architectures
+        is_safe = props.major >= 8  # Allow FP16 only for Ampere (8.x) and later
         _fp16_safe_cache[device] = is_safe
+        if DEBUG_ENABLED:
+            logging.debug(f"FP16 safety check for {props.name}: major={props.major}, is_safe={is_safe}")
         return is_safe
     except Exception:
         _fp16_safe_cache[device] = False
@@ -72,7 +76,8 @@ def clear_vram(device, threshold=0.5, min_free=1.5):
         mem_total = torch.cuda.get_device_properties(device).total_memory / 1024**3
         critical_threshold = 0.05 * mem_total + 0.1  # 5% VRAM + 100 MB
         if mem_allocated > threshold * mem_total or (mem_total - mem_allocated) < max(min_free, critical_threshold):
-            logging.debug(f"Clearing VRAM: allocated {mem_allocated:.2f} GB, free {mem_total - mem_allocated:.2f} GB, threshold {critical_threshold:.2f} GB")
+            if PROFILING_ENABLED:
+                logging.debug(f"Clearing VRAM: allocated {mem_allocated:.2f} GB, free {mem_total - mem_allocated:.2f} GB, threshold {critical_threshold:.2f} GB")
             torch.cuda.empty_cache()
             #soft_empty_cache(clear=False)
             mem_after = torch.cuda.memory_allocated(device) / 1024**3
@@ -362,9 +367,9 @@ def fast_vae_decode(vae, samples):
 def fast_vae_tiled_decode(vae, samples, tile_size=512, overlap=64, temporal_size=64, temporal_overlap=8):
     """Fast VAE decoding with tiling for low VRAM, consistent with fast_vae_decode."""
     device, dtype, is_gpu = initialize_device_and_dtype(vae)
-    vae_dtype = vae_dtype(device=device)
+    vae_dtype_val = vae_dtype(device=device)
     if DEBUG_ENABLED:
-        logging.debug(f"VAE dtype: {vae_dtype}")
+        logging.debug(f"VAE dtype: {vae_dtype_val}")
         logging.debug(f"Pre-VAE checkpoint: {time.time()}")
 
     try:
@@ -378,7 +383,7 @@ def fast_vae_tiled_decode(vae, samples, tile_size=512, overlap=64, temporal_size
             mem_allocated = torch.cuda.memory_allocated(device) / 1024**3
             free_mem = mem_total - mem_allocated
             # Estimate memory for tiled decoding (conservative, ~50% of full decode)
-            vae_memory_required = (vae.memory_used_decode(samples["samples"].shape, vae_dtype) / 1024**3 * 0.5
+            vae_memory_required = (vae.memory_used_decode(samples["samples"].shape, vae_dtype_val) / 1024**3 * 0.5
                                   if hasattr(vae, 'memory_used_decode') else 0.75)
             if PROFILING_ENABLED:
                 logging.debug(f"VRAM before tiled VAE: {mem_allocated:.2f} GB / {mem_total:.2f} GB")
@@ -408,7 +413,7 @@ def fast_vae_tiled_decode(vae, samples, tile_size=512, overlap=64, temporal_size
             latent_samples = samples["samples"]
             if PROFILING_ENABLED:
                 logging.debug(f"Latent samples device: {latent_samples.device}, dtype: {latent_samples.dtype}")
-            latent_samples = optimized_transfer(latent_samples, device, vae_dtype)
+            latent_samples = optimized_transfer(latent_samples, device, vae_dtype_val)
             if is_gpu and force_channels_last():
                 latent_samples = latent_samples.to(memory_format=torch.channels_last)
                 vae.first_stage_model.to(memory_format=torch.channels_last)