iGavroche-rocm-ninodes/nodes.py

#!/usr/bin/env python3
"""
SIMPLE ROCM NODES - EMERGENCY PERFORMANCE FIX
Removed all optimization overhead that was causing slowdown
"""
import torch
import torch.nn.functional as F
import comfy.model_management as model_management
import comfy.utils
import time
import logging
from typing import Dict, Any, Tuple, Optional

class ROCMOptimizedVAEDecode:
    """
    Simple ROCM VAE Decode - No optimization overhead
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "samples": ("LATENT", {"tooltip": "The latent to be decoded."}),
                "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."}),
                "tile_size": ("INT", {
                    "default": 768, 
                    "min": 256, 
                    "max": 2048, 
                    "step": 64,
                    "tooltip": "Tile size for decoding."
                }),
                "overlap": ("INT", {
                    "default": 96, 
                    "min": 32, 
                    "max": 512, 
                    "step": 16,
                    "tooltip": "Overlap between tiles."
                }),
                "use_rocm_optimizations": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable ROCm-specific optimizations."
                }),
                "precision_mode": (["auto", "fp32", "fp16", "mixed"], {
                    "default": "auto",
                    "tooltip": "Precision mode."
                }),
                "batch_optimization": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable batch processing optimizations."
                }),
                "video_chunk_size": ("INT", {
                    "default": 8, 
                    "min": 1, 
                    "max": 32, 
                    "step": 1,
                    "tooltip": "Video chunk size."
                }),
                "memory_optimization_enabled": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable memory optimizations."
                }),
                "adaptive_tiling": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Enable adaptive tiling."
                })
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("IMAGE",)
    FUNCTION = "decode"
    CATEGORY = "ROCM Optimized/VAE"
    
    def decode(self, vae, samples, tile_size=768, overlap=96, use_rocm_optimizations=True, 
               precision_mode="auto", batch_optimization=True, video_chunk_size=8, 
               memory_optimization_enabled=True, adaptive_tiling=True):
        """
        Simple VAE decode - No optimization overhead
        """
        try:
            samples_processed = samples["samples"]
            
            # Simple precision handling
            if precision_mode == "auto":
                optimal_dtype = vae.vae_dtype
            elif precision_mode == "fp16":
                optimal_dtype = torch.float16
            elif precision_mode == "fp32":
                optimal_dtype = torch.float32
            else:
                optimal_dtype = vae.vae_dtype
            
            # Calculate memory requirements
            memory_used = vae.memory_used_decode(samples_processed.shape, optimal_dtype)
            
            # Load VAE
            model_management.load_models_gpu([vae.patcher], memory_required=memory_used)
            
            # Check if we can do direct decode
            if len(samples_processed.shape) == 4:  # Image
                B, C, H, W = samples_processed.shape
                
                # Try direct decode first
                try:
                    with torch.amp.autocast('cuda', enabled=(optimal_dtype != torch.float32), dtype=optimal_dtype):
                        out = vae.first_stage_model.decode(samples_processed.to(optimal_dtype))
                    
                    pixel_samples = vae.process_output(out)
                    pixel_samples = pixel_samples.to(vae.output_device).float()
                    
                    return (pixel_samples,)
                    
                except Exception as e:
                    logging.warning(f"Direct decode failed, falling back to tiled: {e}")
            
            # Fallback to tiled decode
            try:
                pixel_samples = self._decode_tiled(vae, samples, tile_size, overlap)
                return (pixel_samples,)
                
            except Exception as e:
                logging.warning(f"Tiled decode failed, falling back to standard VAE: {e}")
                
                # Final fallback to standard VAE decode
                pixel_samples = vae.decode(samples_processed)
                pixel_samples = vae.process_output(pixel_samples)
                
                return (pixel_samples,)
                
        except Exception as e:
            logging.error(f"VAE decode failed: {e}")
            raise e
    
    def _decode_tiled(self, vae, samples, tile_size, overlap):
        """Simple tiled decode"""
        samples_processed = samples["samples"]
        
        # Handle video tensors (5D) vs image tensors (4D)
        if len(samples_processed.shape) == 5:  # Video: B, C, T, H, W
            B, C, T, H, W = samples_processed.shape
            # For video, process each frame separately
            frames = []
            for t in range(T):
                frame_samples = samples_processed[:, :, t, :, :]  # Remove temporal dimension for processing
                frame_output = self._decode_single_frame(vae, frame_samples, tile_size, overlap)
                frames.append(frame_output)
            return torch.stack(frames, dim=2)  # Stack along temporal dimension
        else:  # Image: B, C, H, W
            return self._decode_single_frame(vae, samples_processed, tile_size, overlap)
    
    def _decode_single_frame(self, vae, samples_processed, tile_size, overlap):
        """Decode a single frame (4D tensor)"""
        # Get compression ratios
        try:
            compression = vae.spacial_compression_decode()
            temporal_compression = vae.temporal_compression_decode()
        except:
            compression = 8
            temporal_compression = 1
        
        # Calculate tile dimensions
        tile_h = tile_size // compression
        tile_w = tile_size // compression
        
        # Calculate number of tiles needed
        B, C, H, W = samples_processed.shape
        tiles_h = (H + tile_h - 1) // tile_h
        tiles_w = (W + tile_w - 1) // tile_w
        
        # Pre-allocate output tensor
        output_shape = (B, 3, H * compression, W * compression)
        pixel_samples = torch.empty(output_shape, dtype=torch.float32, device=samples_processed.device)
        
        # Process tiles
        for i in range(tiles_h):
            for j in range(tiles_w):
                # Calculate tile boundaries with overlap
                start_h = i * tile_h
                end_h = min(start_h + tile_h + overlap, H)
                start_w = j * tile_w
                end_w = min(start_w + tile_w + overlap, W)
                
                # Extract tile
                tile_samples = samples_processed[:, :, start_h:end_h, start_w:end_w]
                
                # Decode tile
                with torch.amp.autocast('cuda', enabled=True, dtype=vae.vae_dtype):
                    tile_output = vae.first_stage_model.decode(tile_samples.to(vae.vae_dtype))
                
                # Process output
                tile_output = vae.process_output(tile_output)
                
                # Copy to output with proper positioning
                output_start_h = start_h * compression
                output_end_h = end_h * compression
                output_start_w = start_w * compression
                output_end_w = end_w * compression
                
                pixel_samples[:, :, output_start_h:output_end_h, output_start_w:output_end_w] = tile_output
        
        return pixel_samples

# Apply no-op instrumentation
def instrument_node(node_class):
    """No-op decorator - instrumentation disabled"""
    return node_class

# Create no-op instrumentation object
class NodeInstrumentation:
    def __init__(self, *args, **kwargs):
        pass
    
    def capture_inputs(self, *args, **kwargs):
        return ""
    
    def capture_outputs(self, *args, **kwargs):
        return ""
    
    def capture_performance(self, *args, **kwargs):
        return ""

# Create global instrumentation object
instrumentation = NodeInstrumentation()

# Apply instrumentation (no-op)
@instrument_node
class ROCMOptimizedVAEDecodeInstrumented(ROCMOptimizedVAEDecode):
    pass

# Node class mapping for ComfyUI
NODE_CLASS_MAPPINGS = {
    "ROCMOptimizedVAEDecode": ROCMOptimizedVAEDecodeInstrumented
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ROCMOptimizedVAEDecode": "ROCM Optimized VAE Decode"
}