hidreamsampler.py

# -*- coding: utf-8 -*-
# HiDream Sampler Node for ComfyUI
# Version: 2024-07-29d (Removed auto-gptq dependency check)
#
# Required Dependencies:
# - transformers, diffusers, torch, numpy, Pillow
# - For NF4 (GPTQ) models: optimum, accelerate (`pip install optimum accelerate`)
#   * Note: Optimum might require additional backends like exllama (`pip install optimum[exllama]`) depending on your setup.
# - For non-NF4/FP8 models (4-bit): bitsandbytes (`pip install bitsandbytes`)
# - Ensure hi_diffusers library is locally available or hdi1 package is installed.
import torch
import numpy as np
from PIL import Image
import comfy.model_management # Ensure this is imported
import comfy.utils
import gc
import os # For checking paths if needed
import huggingface_hub
import importlib.util
from safetensors.torch import load_file

# --- Optional Dependency Handling ---
try:
    import accelerate
    accelerate_available = True
except ImportError:
    accelerate_available = False
    print("Warning: accelerate not installed. device_map='auto' for GPTQ models may not work optimally.")
try:
    import gptqmodel
    gptqmodel_available = True
except ImportError:
    gptqmodel_available = False
    print("Warning: GPTQModel not installed.")
    # Note: Optimum might still load GPTQ without GPTQModel if using ExLlama kernels,
    # but it's often required. Add a warning if NF4 models are selected later.^
try:
    import optimum
    optimum_available = True
    print("Optimum library found. GPTQ model loading enabled (requires suitable backend).")
except ImportError:
    optimum_available = False
    print("Warning: optimum not installed. GPTQ models (NF4 variants) will be disabled.")

try:
    # Import specific classes to avoid potential namespace conflicts later
    from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
    from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
    bnb_available = True
except ImportError:
    bnb_available = False
    # Keep placeholders None to avoid errors later if bnb not available
    TransformersBitsAndBytesConfig = None
    DiffusersBitsAndBytesConfig = None
    print("Warning: bitsandbytes not installed. 4-bit BNB quantization will not be available.")

# --- Core Imports ---
from transformers import LlamaForCausalLM, AutoTokenizer # Use AutoTokenizer

# --- HiDream Specific Imports ---
# Attempt local import first, then fallback (which might fail)
try:
    # Assuming hi_diffusers is cloned into this custom_node's directory
    from .hi_diffusers.models.transformers.transformer_hidream_image import HiDreamImageTransformer2DModel
    from .hi_diffusers.pipelines.hidream_image.pipeline_hidream_image import HiDreamImagePipeline
    from .hi_diffusers.pipelines.hidream_image.pipeline_hidream_image_to_image import HiDreamImageToImagePipeline
    from .hi_diffusers.schedulers.fm_solvers_unipc import FlowUniPCMultistepScheduler
    from .hi_diffusers.schedulers.flash_flow_match import FlashFlowMatchEulerDiscreteScheduler
    hidream_classes_loaded = True
except ImportError as e:
    print("--------------------------------------------------------------------")
    print(f"ComfyUI-HiDream-Sampler: Could not import local hi_diffusers ({e}).")
    print("Please ensure hi_diffusers library is inside ComfyUI-HiDream-Sampler,")
    print("or hdi1 package is installed in the ComfyUI environment.")
    print("Node may fail to load models.")
    print("--------------------------------------------------------------------")
    # Define placeholders so the script doesn't crash immediately
    HiDreamImageTransformer2DModel = None
    HiDreamImagePipeline = None
    FlowUniPCMultistepScheduler = None
    FlashFlowMatchEulerDiscreteScheduler = None
    hidream_classes_loaded = False

# --- Model Paths ---
ORIGINAL_MODEL_PREFIX = "HiDream-ai"
NF4_MODEL_PREFIX = "azaneko"

ORIGINAL_LLAMA_MODEL_NAME = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" 
NF4_LLAMA_MODEL_NAME = "ModelCloud/Meta-Llama-3.1-8B-Instruct-gptq-4bit"
UNCENSORED_LLAMA_MODEL_NAME = "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2" 
UNCENSORED_NF4_LLAMA_MODEL_NAME = "shuttercat/DarkIdol-Llama3.1-NF4-GPTQ"

# --- Model Configurations ---
MODEL_CONFIGS = {
    # --- NF4 Models (Require Optimum) ---
    "full-nf4": {
        "path": f"{NF4_MODEL_PREFIX}/HiDream-I1-Full-nf4",
        "guidance_scale": 5.0, "num_inference_steps": 50, "shift": 3.0,
        "scheduler_class": "FlowUniPCMultistepScheduler", 
        "is_nf4": True, "is_fp8": False, "requires_bnb": False, "requires_gptq_deps": True # Requires Optimum
    },
    "dev-nf4": {
        "path": f"{NF4_MODEL_PREFIX}/HiDream-I1-Dev-nf4",
        "guidance_scale": 0.0, "num_inference_steps": 28, "shift": 6.0,
        "scheduler_class": "FlashFlowMatchEulerDiscreteScheduler",
        "is_nf4": True, "is_fp8": False, "requires_bnb": False, "requires_gptq_deps": True # Requires Optimum
    },
    "fast-nf4": {
        "path": f"{NF4_MODEL_PREFIX}/HiDream-I1-Fast-nf4",
        "guidance_scale": 0.0, "num_inference_steps": 16, "shift": 3.0,
        "scheduler_class": "FlashFlowMatchEulerDiscreteScheduler",
        "is_nf4": True, "is_fp8": False, "requires_bnb": False, "requires_gptq_deps": True # Requires Optimum
    },
    # --- Original/BNB Models (Require BitsAndBytes) ---
     "full": {
        "path": f"{ORIGINAL_MODEL_PREFIX}/HiDream-I1-Full",
        "guidance_scale": 5.0, "num_inference_steps": 50, "shift": 3.0,
        "scheduler_class": "FlowUniPCMultistepScheduler",
        "is_nf4": False, "is_fp8": False, "requires_bnb": True, "requires_gptq_deps": False
    },
    "dev": {
        "path": f"{ORIGINAL_MODEL_PREFIX}/HiDream-I1-Dev",
        "guidance_scale": 0.0, "num_inference_steps": 28, "shift": 6.0,
        "scheduler_class": "FlashFlowMatchEulerDiscreteScheduler",
        "is_nf4": False, "is_fp8": False, "requires_bnb": True, "requires_gptq_deps": False
    },
    "fast": {
        "path": f"{ORIGINAL_MODEL_PREFIX}/HiDream-I1-Fast",
        "guidance_scale": 0.0, "num_inference_steps": 16, "shift": 3.0,
        "scheduler_class": "FlashFlowMatchEulerDiscreteScheduler",
        "is_nf4": False, "is_fp8": False, "requires_bnb": True, "requires_gptq_deps": False
    }
}

# --- Filter models based on available dependencies ---
# (Keep filtering logic the same)
original_model_count = len(MODEL_CONFIGS)
if not bnb_available: MODEL_CONFIGS = {k: v for k, v in MODEL_CONFIGS.items() if not v.get("requires_bnb", False)}
if not optimum_available or not gptqmodel_available: MODEL_CONFIGS = {k: v for k, v in MODEL_CONFIGS.items() if not v.get("requires_gptq_deps", False)}
if not hidream_classes_loaded: MODEL_CONFIGS = {}
filtered_model_count = len(MODEL_CONFIGS)
if filtered_model_count == 0: print("*"*70 + "\nCRITICAL ERROR: No HiDream models available...\n" + "*"*70)
elif filtered_model_count < original_model_count: print("*"*70 + "\nWarning: Some HiDream models disabled...\n" + "*"*70)
# Define BitsAndBytes configs (if available)
# (Keep definitions the same)
bnb_llm_config = None; bnb_transformer_4bit_config = None
if bnb_available: bnb_llm_config = TransformersBitsAndBytesConfig(load_in_4bit=True); bnb_transformer_4bit_config = DiffusersBitsAndBytesConfig(load_in_4bit=True)
model_dtype = torch.bfloat16
# Get available scheduler classes
# (Keep definitions the same)
available_schedulers = {};
if hidream_classes_loaded: available_schedulers = {"FlowUniPCMultistepScheduler": FlowUniPCMultistepScheduler, "FlashFlowMatchEulerDiscreteScheduler": FlashFlowMatchEulerDiscreteScheduler}

DEBUG_CACHE = True  # Set to False in production

# Use a more aggressive global cleanup
def global_cleanup():
    """Global cleanup function for use with multiple HiDream nodes"""
    print("HiDream: Performing global cleanup...")
    
    # Clear any pending operations
    torch.cuda.synchronize()
    
    # Get current memory stats
    if torch.cuda.is_available():
        before_mem = torch.cuda.memory_allocated() / 1024**2
        print(f"  Memory before cleanup: {before_mem:.2f} MB")
    
    # Perform HiDreamSampler cleanup
    HiDreamSampler.cleanup_models()
    
    # Additional cleanup
    gc.collect()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        after_mem = torch.cuda.memory_allocated() / 1024**2
        print(f"  Memory after cleanup: {after_mem:.2f} MB")
    
    return True

# --- Helper: Get Scheduler Instance ---
def get_scheduler_instance(scheduler_name, shift_value):
    if not available_schedulers: raise RuntimeError("No schedulers available...")
    scheduler_class = available_schedulers.get(scheduler_name)
    if scheduler_class is None: raise ValueError(f"Scheduler class '{scheduler_name}' not found...")
    return scheduler_class(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
# --- Loading Function (Handles NF4 and default BNB) --- 
def load_models(model_type, use_uncensored_llm=False):
    if not hidream_classes_loaded: raise ImportError("Cannot load models: HiDream classes failed to import.")
    if model_type not in MODEL_CONFIGS: raise ValueError(f"Unknown or incompatible model_type: {model_type}")
    config = MODEL_CONFIGS[model_type]
    model_path = config["path"]; is_nf4 = config.get("is_nf4", False)
    scheduler_name = config["scheduler_class"]; shift = config["shift"]
    requires_bnb = config.get("requires_bnb", False); requires_gptq_deps = config.get("requires_gptq_deps", False)
    if requires_bnb and not bnb_available: raise ImportError(f"Model '{model_type}' requires BitsAndBytes...")
    if requires_gptq_deps and (not optimum_available or not gptqmodel_available): raise ImportError(f"Model '{model_type}' requires Optimum & AutoGPTQ...")
    print(f"--- Loading Model Type: {model_type} ---"); print(f"Model Path: {model_path}")
    print(f"NF4: {is_nf4}, Requires BNB: {requires_bnb}, Requires GPTQ deps: {requires_gptq_deps}")
    print(f"Using Uncensored LLM: {use_uncensored_llm}")
    start_mem = torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0; print(f"(Start VRAM: {start_mem:.2f} MB)")

    # Create a standardized cache key used by all nodes
    cache_key = f"{model_type}_{'uncensored' if use_uncensored_llm else 'standard'}"
    
    # Check cache with debug info
    if DEBUG_CACHE:
        print(f"Cache check for key: {cache_key}")
        print(f"Cache contains: {list(HiDreamSampler._model_cache.keys())}")
    
    if cache_key in HiDreamSampler._model_cache:
        pipe, stored_config = HiDreamSampler._model_cache[cache_key]
        if pipe is not None and hasattr(pipe, 'transformer') and pipe.transformer is not None:
            print(f"Using cached model for {cache_key}")
            return pipe, MODEL_CONFIGS[model_type]  # Always return original config dict
        else:
            print(f"Cache entry invalid for {cache_key}, reloading")
            # Remove from cache to avoid reusing
            HiDreamSampler._model_cache.pop(cache_key, None)

    
    # --- 1. Load LLM (Conditional) ---
    text_encoder_load_kwargs = {"low_cpu_mem_usage": True, "torch_dtype": model_dtype}
    
    if is_nf4:
        # Choose uncensored model if requested, but keep loading process identical
        if use_uncensored_llm:
            llama_model_name = UNCENSORED_NF4_LLAMA_MODEL_NAME
            print(f"\n[1a] Preparing Uncensored LLM (GPTQ): {llama_model_name}")
        else:
            llama_model_name = NF4_LLAMA_MODEL_NAME
            print(f"\n[1a] Preparing LLM (GPTQ): {llama_model_name}")
            
        # Rest of the NF4 loading process stays exactly the same
        if accelerate_available:
            # Fix for device format - use integer instead of cuda:0
            if hasattr(torch.cuda, 'get_device_properties') and torch.cuda.is_available():
                total_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
                # Use 40% for model, leaving room for the transformer
                max_mem = int(total_mem * 0.4)
                text_encoder_load_kwargs["max_memory"] = {0: f"{max_mem}GiB"}
                print(f"     Setting max memory limit: {max_mem}GiB of {total_mem:.1f}GiB")
            text_encoder_load_kwargs["device_map"] = "auto";
            print("     Using device_map='auto'.")
        else: print("     accelerate not found, attempting manual placement.")
    else:
        # For non-NF4 models, choose uncensored if requested
        if use_uncensored_llm:
            llama_model_name = UNCENSORED_LLAMA_MODEL_NAME
            print(f"\n[1a] Preparing Uncensored LLM (4-bit BNB): {llama_model_name}")
        else:
            llama_model_name = ORIGINAL_LLAMA_MODEL_NAME
            print(f"\n[1a] Preparing LLM (4-bit BNB): {llama_model_name}")
            
        # Rest of standard model loading stays exactly the same
        if bnb_llm_config: text_encoder_load_kwargs["quantization_config"] = bnb_llm_config; print("     Using 4-bit BNB.")
        else: raise ImportError("BNB config required for standard LLM.")
    
    print(f"[1b] Loading Tokenizer: {llama_model_name}..."); tokenizer = AutoTokenizer.from_pretrained(llama_model_name, use_fast=False); print("     Tokenizer loaded.")

    if is_nf4:
        # More aggressive RoPE scaling fix
        try:
            # Direct fix: Load and completely replace the config
            from transformers import AutoConfig
            
            # Load the config and make a deep copy we can modify
            config = AutoConfig.from_pretrained(llama_model_name)
            
            """
            # Completely replace the rope_scaling with a known good config
            config.rope_scaling = {"type": "linear", "factor": 1.0}
            print(f"     ✅ Fixed rope_scaling to: {config.rope_scaling}") 
            """
            
            # Force using our fixed config
            text_encoder_load_kwargs["config"] = config
            
            # Disable all validation for good measure
            text_encoder_load_kwargs["low_cpu_mem_usage"] = True
        except Exception as e:
            print(f"     ⚠️ Failed to patch config: {e}")
    
    print(f"[1c] Loading Text Encoder: {llama_model_name}... (May download files)"); text_encoder = LlamaForCausalLM.from_pretrained(llama_model_name, **text_encoder_load_kwargs)
    if "device_map" not in text_encoder_load_kwargs: print("     Moving text encoder to CUDA..."); text_encoder.to("cuda")
    step1_mem = torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0; print(f"✅ Text encoder loaded! (VRAM: {step1_mem:.2f} MB)")
    
    # --- 2. Load Transformer (Conditional) ---
    print(f"\n[2] Preparing Transformer from: {model_path}"); transformer_load_kwargs = {"subfolder": "transformer", "torch_dtype": model_dtype, "low_cpu_mem_usage": True}
    if is_nf4: print("     Type: NF4")
    else: # Default BNB case
        print("     Type: Standard (Applying 4-bit BNB quantization)")
        if bnb_transformer_4bit_config:
            transformer_load_kwargs["quantization_config"] = bnb_transformer_4bit_config
        else:
            raise ImportError("BNB config required for transformer but unavailable.")
    print("     Loading Transformer... (May download files)"); transformer = HiDreamImageTransformer2DModel.from_pretrained(model_path, **transformer_load_kwargs)
    print("     Moving Transformer to CUDA..."); transformer.to("cuda")
    step2_mem = torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0; print(f"✅ Transformer loaded! (VRAM: {step2_mem:.2f} MB)")
    
    # --- 3. Load Scheduler ---
    # Load the scheduler using the default shift value from the config initially
    print(f"\n[3] Preparing Scheduler: {scheduler_name} (Default shift: {shift})")
    scheduler = get_scheduler_instance(scheduler_name, shift)
    print(f"     Using Scheduler: {scheduler_name}")

    # --- 4. Load Pipeline ---
    print(f"\n[4] Loading Pipeline from: {model_path}"); print("     Passing pre-loaded components...")
    pipe = HiDreamImagePipeline.from_pretrained(model_path, scheduler=scheduler, tokenizer_4=tokenizer, text_encoder_4=text_encoder, transformer=None, torch_dtype=model_dtype, low_cpu_mem_usage=True); print("     Pipeline structure loaded.")
    
    # --- 5. Final Setup ---
    print("\n[5] Finalizing Pipeline..."); print("     Assigning transformer..."); pipe.transformer = transformer
    print("     Moving pipeline object to CUDA (final check)...");
    try: pipe.to("cuda")
    except Exception as e: print(f"     Warning: Could not move pipeline object to CUDA: {e}.")
    if is_nf4:
        print("     Attempting CPU offload for NF4...");
        if hasattr(pipe, "enable_sequential_cpu_offload"):
            try: pipe.enable_sequential_cpu_offload(); print("     ✅ CPU offload enabled.")
            except Exception as e: print(f"     ⚠️ Failed CPU offload: {e}")
        else: print("     ⚠️ enable_sequential_cpu_offload() not found.")
    final_mem = torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0; print(f"✅ Pipeline ready! (VRAM: {final_mem:.2f} MB)")
    return pipe, MODEL_CONFIGS[model_type]
    
# --- Resolution Parsing & Tensor Conversion ---
RESOLUTION_OPTIONS = [ # (Keep list the same)
    "1024 × 1024 (Square)","768 × 1360 (Portrait)","1360 × 768 (Landscape)",
    "880 × 1168 (Portrait)","1168 × 880 (Landscape)","1248 × 832 (Landscape)",
    "832 × 1248 (Portrait)"
]

def parse_resolution(resolution_str):
    """Parse resolution string into height and width dimensions."""
    try:
        # Extract the resolution part before the parenthesis
        res_part = resolution_str.split(" (")[0].strip()
        # Replace 'x' with '×' for consistency if needed
        parts = res_part.replace('x', '×').split("×")
        
        if len(parts) != 2:
            raise ValueError(f"Expected format 'width × height', got '{res_part}'")
            
        width_str = parts[0].strip()
        height_str = parts[1].strip()
        
        width = int(width_str)
        height = int(height_str)
        print(f"Successfully parsed resolution: {width}x{height}")
        return height, width
    except Exception as e:
        print(f"Error parsing resolution '{resolution_str}': {e}. Falling back to 1024x1024.")
        return 1024, 1024
    
def pil2tensor(image: Image.Image):
    """Convert PIL image to tensor with better error handling"""
    if image is None:
        print("pil2tensor: Image is None")
        return None

    try:
        # Debug image properties
        print(f"pil2tensor: Image mode={image.mode}, size={image.size}")
        
        # Ensure image is in RGB mode
        if image.mode != 'RGB':
            print(f"Converting image from {image.mode} to RGB")
            image = image.convert('RGB')
        
        # Convert to numpy array with explicit steps
        np_array = np.array(image)
        print(f"Numpy array shape={np_array.shape}, dtype={np_array.dtype}")
        
        # Convert to float32 and normalize
        np_array = np_array.astype(np.float32) / 255.0
        
        # Convert to tensor and add batch dimension
        tensor = torch.from_numpy(np_array)
        tensor = tensor.unsqueeze(0)
        print(f"Final tensor shape={tensor.shape}")
        
        return tensor
    except Exception as e:
        print(f"Error in pil2tensor: {e}")
        import traceback
        traceback.print_exc()
        
        # Try ComfyUI's own conversion if ours fails
        try:
            print("Trying ComfyUI's own conversion...")
            tensor = comfy.utils.pil2tensor(image)
            print(f"ComfyUI conversion successful: {tensor.shape}")
            return tensor
        except Exception as e2:
            print(f"ComfyUI conversion also failed: {e2}")
            return None

# --- ComfyUI Node Definition ---
class HiDreamSampler:
    _model_cache = {}
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate"
    CATEGORY = "HiDream"
    
    @classmethod
    def cleanup_models(cls):
        """Clean up all cached models - can be called by external memory management"""
        print("HiDream: Cleaning up all cached models...")
        keys_to_del = list(cls._model_cache.keys())
        for key in keys_to_del:
            print(f"  Removing '{key}'...")
            try:
                pipe_to_del, _ = cls._model_cache.pop(key)
                # More aggressive cleanup - clear all major components
                if hasattr(pipe_to_del, 'transformer'):
                    pipe_to_del.transformer = None
                if hasattr(pipe_to_del, 'text_encoder_4'):
                    pipe_to_del.text_encoder_4 = None
                if hasattr(pipe_to_del, 'tokenizer_4'):
                    pipe_to_del.tokenizer_4 = None
                if hasattr(pipe_to_del, 'scheduler'):
                    pipe_to_del.scheduler = None
                del pipe_to_del
            except Exception as e:
                print(f"  Error cleaning up {key}: {e}")
        # Multiple garbage collection passes
        for _ in range(3):
            gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            # Force synchronization
            torch.cuda.synchronize()
        print("HiDream: Cache cleared")
        return True
        
    @staticmethod
    def parse_aspect_ratio(aspect_ratio_str):
        """Parse aspect ratio string to get width and height"""
        try:
            # Extract dimensions from the parenthesis
            dims_part = aspect_ratio_str.split("(")[1].split(")")[0]
            width, height = dims_part.split("×")
            return int(width), int(height)
        except Exception as e:
            print(f"Error parsing aspect ratio '{aspect_ratio_str}': {e}. Falling back to 1024x1024.")
            return 1024, 1024
            
    @classmethod
    def INPUT_TYPES(s):
        available_model_types = list(MODEL_CONFIGS.keys())
        if not available_model_types:
            return {"required": {"error": ("STRING", {"default": "No models available...", "multiline": True})}}
        default_model = "fast-nf4" if "fast-nf4" in available_model_types else "fast" if "fast" in available_model_types else available_model_types[0]
        
        # Define schedulers
        scheduler_options = [
            "Default for model",
            "UniPC",
            "Euler",
            "Karras Euler",
            "Karras Exponential"
        ]
        
        return {
            "required": {
                "model_type": (available_model_types, {"default": default_model}),
                "prompt": ("STRING", {"multiline": True, "default": "..."}),
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "resolution": (RESOLUTION_OPTIONS, {"default": "1024 × 1024 (Square)"}),
                "num_images": ("INT", {"default": 1, "min": 0, "max": 8,}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "scheduler": (scheduler_options, {"default": "Default for model"}),
                "override_steps": ("INT", {"default": -1, "min": -1, "max": 100}),
                "override_cfg": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}),
                "override_shift": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}), # <-- Added Shift Override
                "override_width": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                "override_height": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8})
            }
        }
    
    def generate(self, model_type, prompt, negative_prompt, resolution, num_images, seed, scheduler, override_steps, override_cfg, override_shift, override_width, override_height): # <-- Added override_shift
        use_uncensored_llm = None

        # Determine resolution
        if override_width > 0 and override_height > 0:
            height, width = override_height, override_width
            print(f"Using override resolution: {width}x{height}")
        else:
            height, width = parse_resolution(resolution)
            print(f"Using fixed resolution: {width}x{height} ({resolution})")

        # Monitor initial memory usage
        if torch.cuda.is_available():
            initial_mem = torch.cuda.memory_allocated() / 1024**2
            print(f"HiDream: Initial VRAM usage: {initial_mem:.2f} MB")
        if not MODEL_CONFIGS or model_type == "error":
            print("HiDream Error: No models loaded.")
            return (torch.zeros((1, 512, 512, 3)),)
        
        pipe = None; config = None
        cache_key = f"{model_type}"
        
        # --- Model Loading / Caching ---
        if cache_key in self._model_cache:
            print(f"Checking cache for {cache_key}...")
            pipe, config = self._model_cache[cache_key]
            valid_cache = True
            if pipe is None or config is None or not hasattr(pipe, 'transformer') or pipe.transformer is None:
                valid_cache = False
                print("Invalid cache, reloading...")
                del self._model_cache[cache_key]
                pipe, config = None, None
            if valid_cache:
                print("Using cached model.")

        if pipe is None:
            if self._model_cache:
                print(f"Clearing ALL cache before loading {model_type}...")
                keys_to_del = list(self._model_cache.keys())
                for key in keys_to_del:
                    print(f"  Removing '{key}'...")
                    try:
                        pipe_to_del, _= self._model_cache.pop(key)
                        if hasattr(pipe_to_del, 'transformer'):
                            pipe_to_del.transformer = None
                        if hasattr(pipe_to_del, 'text_encoder_4'):
                            pipe_to_del.text_encoder_4 = None
                        if hasattr(pipe_to_del, 'tokenizer_4'):
                            pipe_to_del.tokenizer_4 = None
                        if hasattr(pipe_to_del, 'scheduler'):
                            pipe_to_del.scheduler = None
                        del pipe_to_del
                    except Exception as e:
                        print(f"  Error removing {key}: {e}")

                    # Multiple garbage collection passes
                    for _ in range(3):
                        gc.collect()
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                        # Force synchronization
                        torch.cuda.synchronize()
                    print("Cache cleared.")

            print(f"Loading model for {model_type}...")

            try:
                pipe, config = load_models(model_type, use_uncensored_llm)
                self._model_cache[cache_key] = (pipe, config)
                print(f"Model {model_type}{' (uncensored)' if use_uncensored_llm else ''} loaded & cached!")
            except Exception as e:
                print(f"!!! ERROR loading {model_type}: {e}")
                import traceback
                traceback.print_exc()
                return (torch.zeros((1, 512, 512, 3)),)
        
        if pipe is None or config is None:
            print("CRITICAL ERROR: Load failed.")
            return (torch.zeros((1, 512, 512, 3)),)
        
        # --- Update scheduler if requested ---
        original_scheduler_class = config["scheduler_class"]
        # Determine shift value: override if >= 0.0, else use config default
        default_shift = config["shift"]
        shift_value = override_shift if override_shift >= 0.0 else default_shift
        print(f"Selected Shift Value: {shift_value} (Override: {override_shift}, Default: {default_shift})")



        if scheduler != "Default for model":
            print(f"Replacing default scheduler ({original_scheduler_class}) with: {scheduler}")

            # Create a completely fresh scheduler instance to avoid any parameter leakage
            if scheduler == "UniPC":
                new_scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000,
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_karras_sigmas=True
                )
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Exponential":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000,
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_exponential_sigmas=True
                )
                pipe.scheduler = new_scheduler

        else:
            # Ensure we're using the original scheduler as specified in the model config
            print(f"Using model's default scheduler type: {original_scheduler_class} with shift={shift_value}")
            pipe.scheduler = get_scheduler_instance(original_scheduler_class, shift_value)

        # --- Generation Setup ---
        is_nf4_current = config.get("is_nf4", False)
        num_inference_steps = override_steps if override_steps >= 0 else config["num_inference_steps"]
        guidance_scale = override_cfg if override_cfg >= 0.0 else config["guidance_scale"]

        # Create the progress bar
        pbar = comfy.utils.ProgressBar(num_inference_steps)

        # Set default max sequence lengths
        max_length_clip_l = 77
        max_length_openclip = 150
        max_length_t5 = 256
        max_length_llama = 256

        # Set default encoder weights
        clip_l_weight = 1.0
        openclip_weight = 1.0
        t5_weight = 1.0
        llama_weight = 1.0
                    
        try:
            inference_device = comfy.model_management.get_torch_device()
        except Exception:
            inference_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        print(f"Creating Generator on: {inference_device}")
        generator = torch.Generator(device=inference_device).manual_seed(seed)
        print(f"\n--- Starting Generation ---")
        print(f"Model: {model_type}, Res: {height}x{width}, Steps: {num_inference_steps}, CFG: {guidance_scale}, Shift: {pipe.scheduler.config.shift}, Seed: {seed}")
        print(f"Using standard sequence lengths: CLIP-L: {max_length_clip_l}, OpenCLIP: {max_length_openclip}, T5: {max_length_t5}, Llama: {max_length_llama}")

        # --- Run Inference ---
        pipeline_output = None
        try:
            if not is_nf4_current:
                print(f"Ensuring pipe on: {inference_device} (Offload NOT enabled)")
                pipe.to(inference_device)
            else:
                print(f"Skipping pipe.to({inference_device}) (CPU offload enabled).")
            print("Executing pipeline inference...")

            # Ensure batch size consistency for multiple images
            if num_images > 1:
                print(f"Preparing for batch generation with {num_images} images...")
                # Create a list to store outputs
                output_images_list = []
                for i in range(num_images):
                    print(f"Generating image {i+1}/{num_images}...")
                    # Generate one image at a time to avoid batch size issues
                    with torch.inference_mode():
                        single_output = pipe(
                            prompt=prompt,
                            prompt_2=prompt,
                            prompt_3=prompt,
                            prompt_4=prompt,
                            negative_prompt=negative_prompt.strip() if negative_prompt else None,
                            height=height,
                            width=width,
                            guidance_scale=guidance_scale,
                            num_inference_steps=num_inference_steps,
                            num_images_per_prompt=1,  # Force single image per call
                            generator=torch.Generator(device=inference_device).manual_seed(seed + i),  # Increment seed for variety
                            max_sequence_length_clip_l=max_length_clip_l,
                            max_sequence_length_openclip=max_length_openclip,
                            max_sequence_length_t5=max_length_t5,
                            max_sequence_length_llama=max_length_llama,
                            clip_l_scale=clip_l_weight,
                            openclip_scale=openclip_weight,
                            t5_scale=t5_weight,
                            llama_scale=llama_weight,
                        )
                    output_images_list.extend(single_output.images)
                    pbar.update_absolute((i + 1) * num_inference_steps // num_images)
            else:
                with torch.inference_mode():
                    pipeline_output = pipe(
                        prompt=prompt,
                        prompt_2=prompt,
                        prompt_3=prompt,
                        prompt_4=prompt,
                        negative_prompt=negative_prompt.strip() if negative_prompt else None,
                        height=height,
                        width=width,
                        guidance_scale=guidance_scale,
                        num_inference_steps=num_inference_steps,
                        num_images_per_prompt=num_images,
                        generator=generator,
                        max_sequence_length_clip_l=max_length_clip_l,
                        max_sequence_length_openclip=max_length_openclip,
                        max_sequence_length_t5=max_length_t5,
                        max_sequence_length_llama=max_length_llama,
                        clip_l_scale=clip_l_weight,
                        openclip_scale=openclip_weight,
                        t5_scale=t5_weight,
                        llama_scale=llama_weight,
                    )
                    output_images_list = pipeline_output.images

            print("Pipeline inference finished.")

        except Exception as e:
            print(f"!!! ERROR during execution: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, height, width, 3)),)
        finally:
            pbar.update_absolute(num_inference_steps) # Update pbar regardless
        print("--- Generation Complete ---")
        
        # Robust output handling
        if output_images_list is None or not isinstance(output_images_list, list) or len(output_images_list) == 0:
            print(f"ERROR: No images returned or invalid format (Type: {type(output_images_list)}). Creating blank image.")
            return (torch.zeros((1, height, width, 3)),)

        try:
            print(f"Processing {len(output_images_list)} output image(s).")
            tensor_list = []

            for i, img in enumerate(output_images_list):
                if not isinstance(img, Image.Image):
                    print(f"WARNING: Item {i} in output list is not a PIL Image (Type: {type(img)}). Skipping.")
                    continue

                print(f"Converting image {i+1}/{len(output_images_list)}...")
                single_tensor = pil2tensor(img) # This returns shape [1, H, W, C]

                if single_tensor is not None:
                    if len(single_tensor.shape) == 4 and single_tensor.shape[0] == 1:
                        tensor_list.append(single_tensor)
                    else:
                        print(f"WARNING: pil2tensor returned unexpected shape {single_tensor.shape} for image {i}. Skipping.")
                else:
                    print(f"WARNING: pil2tensor failed for image {i}. Skipping.")

            if not tensor_list:
                print("ERROR: All image conversions failed. Creating blank image.")
                return (torch.zeros((1, height, width, 3)),)

            output_tensor = torch.cat(tensor_list, dim=0)
            print(f"Successfully converted {output_tensor.shape[0]} images into batch tensor.")

            if output_tensor.dtype != torch.float32:
                print(f"Converting batched {output_tensor.dtype} tensor to float32 for ComfyUI compatibility")
                output_tensor = output_tensor.to(torch.float32)

            if len(output_tensor.shape) != 4 or output_tensor.shape[0] == 0 or output_tensor.shape[3] != 3:
                print(f"ERROR: Invalid final batch tensor shape {output_tensor.shape}. Creating blank image.")
                return (torch.zeros((1, height, width, 3)),)

            print(f"Output tensor shape: {output_tensor.shape}") 

            try:
                import comfy.model_management as model_management
                print("HiDream: Requesting ComfyUI memory cleanup...")
                model_management.soft_empty_cache()
            except Exception as e:
                print(f"HiDream: ComfyUI cleanup failed: {e}")

            if torch.cuda.is_available():
                final_mem = torch.cuda.memory_allocated() / 1024**2
                print(f"HiDream: Final VRAM usage: {final_mem:.2f} MB (Change: {final_mem-initial_mem:.2f} MB)")
                
            return (output_tensor,)
            
        except Exception as e:
            print(f"Error processing output image: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, height, width, 3)),)

class HiDreamSamplerAdvanced:
    _model_cache = HiDreamSampler._model_cache
    cleanup_models = HiDreamSampler.cleanup_models
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate"
    CATEGORY = "HiDream"
    
    @classmethod
    def INPUT_TYPES(s):
        available_model_types = list(MODEL_CONFIGS.keys())
        if not available_model_types:
            return {"required": {"error": ("STRING", {"default": "No models available...", "multiline": True})}}
        default_model = "fast-nf4" if "fast-nf4" in available_model_types else "fast" if "fast" in available_model_types else available_model_types[0]
        
        # Define schedulers
        scheduler_options = [
            "Default for model",
            "UniPC",
            "Euler",
            "Karras Euler",
            "Karras Exponential"
        ]
        
        return {
            "required": {
                "model_type": (available_model_types, {"default": default_model}),
                "primary_prompt": ("STRING", {"multiline": True, "default": "..."}),
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "resolution": (RESOLUTION_OPTIONS, {"default": "1024 × 1024 (Square)"}),
                "num_images": ("INT", {"default": 1, "min": 0, "max": 8,}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "scheduler": (scheduler_options, {"default": "Default for model"}),
                "override_steps": ("INT", {"default": -1, "min": -1, "max": 100}),
                "override_cfg": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}),
                "override_shift": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}), # <-- Added Shift Override
                "override_width": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                "override_height": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 8}),
                "use_uncensored_llm": ("BOOLEAN", {"default": False})
            },
            "optional": {
                "clip_l_prompt": ("STRING", {"multiline": True, "default": ""}),
                "openclip_prompt": ("STRING", {"multiline": True, "default": ""}),
                "t5_prompt": ("STRING", {"multiline": True, "default": ""}),
                "llama_prompt": ("STRING", {"multiline": True, "default": ""}),
                "llm_system_prompt": ("STRING", {
                    "multiline": True, 
                    "default": "You are a creative AI assistant that helps create detailed, vivid images based on user descriptions."
                }),
                "clip_l_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "openclip_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "t5_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "llama_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "max_length_clip_l": ("INT", {"default": 77, "min": 64, "max": 218}),
                "max_length_openclip": ("INT", {"default": 77, "min": 64, "max": 218}),
                "max_length_t5": ("INT", {"default": 128, "min": 64, "max": 512}),
                "max_length_llama": ("INT", {"default": 128, "min": 64, "max": 2048})
            }
        }
    
    def generate(self, model_type, primary_prompt, negative_prompt, resolution, num_images, seed, scheduler, override_steps, override_cfg, override_shift, use_uncensored_llm=False, clip_l_prompt="", openclip_prompt="", t5_prompt="", llama_prompt="", llm_system_prompt="You are a creative AI assistant...", override_width=0, override_height=0, max_length_clip_l=77, max_length_openclip=77, max_length_t5=128, max_length_llama=128,clip_l_weight=1.0, openclip_weight=1.0, t5_weight=1.0, llama_weight=1.0, **kwargs): # <-- Added override_shift
        
        # Determine resolution
        if override_width > 0 and override_height > 0:
            height, width = override_height, override_width
            print(f"Using override resolution: {width}x{height}")
        else:
            height, width = parse_resolution(resolution)
            print(f"Using fixed resolution: {width}x{height} ({resolution})")
        
        # Monitor initial memory usage
        if torch.cuda.is_available():
            initial_mem = torch.cuda.memory_allocated() / 1024**2
            print(f"HiDream: Initial VRAM usage: {initial_mem:.2f} MB")
            
        if not MODEL_CONFIGS or model_type == "error":
            print("HiDream Error: No models loaded.")
            return (torch.zeros((1, 512, 512, 3)),)
            
        pipe = None; config = None
        cache_key = f"{model_type}_{'uncensored' if use_uncensored_llm else 'standard'}"
        
        # --- Model Loading / Caching ---
        if cache_key in self._model_cache:
            print(f"Checking cache for {cache_key}...")
            pipe, config = self._model_cache[cache_key]
            valid_cache = True
            if pipe is None or config is None or not hasattr(pipe, 'transformer') or pipe.transformer is None:
                valid_cache = False
                print("Invalid cache, reloading...")
                del self._model_cache[cache_key]
                pipe, config = None, None
            if valid_cache:
                print("Using cached model.")
                
        if pipe is None:
            if self._model_cache:
                print(f"Clearing ALL cache before loading {model_type}...")
                keys_to_del = list(self._model_cache.keys())
                for key in keys_to_del:
                    print(f"  Removing '{key}'...")
                    try:
                        pipe_to_del, _= self._model_cache.pop(key)
                        if hasattr(pipe_to_del, 'transformer'):
                            pipe_to_del.transformer = None
                        if hasattr(pipe_to_del, 'text_encoder_4'):
                            pipe_to_del.text_encoder_4 = None
                        if hasattr(pipe_to_del, 'tokenizer_4'):
                            pipe_to_del.tokenizer_4 = None
                        if hasattr(pipe_to_del, 'scheduler'):
                            pipe_to_del.scheduler = None
                        del pipe_to_del
                    except Exception as e:
                        print(f"  Error removing {key}: {e}")

                # Multiple garbage collection passes
                for _ in range(3):
                    gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    # Force synchronization
                    torch.cuda.synchronize()
                print("Cache cleared.")
                
            print(f"Loading model for {model_type}{' (uncensored)' if use_uncensored_llm else ''}...")

            try:
                pipe, config = load_models(model_type, use_uncensored_llm)
                self._model_cache[cache_key] = (pipe, config)
                print(f"Model {model_type}{' (uncensored)' if use_uncensored_llm else ''} loaded & cached!")
            except Exception as e:
                print(f"!!! ERROR loading {model_type}: {e}")
                import traceback
                traceback.print_exc()
                return (torch.zeros((1, 512, 512, 3)),)
                
        if pipe is None or config is None:
            print("CRITICAL ERROR: Load failed.")
            return (torch.zeros((1, 512, 512, 3)),)
            
        # --- Update scheduler if requested ---
        original_scheduler_class = config["scheduler_class"]
        # Determine shift value: override if >= 0.0, else use config default
        default_shift = config["shift"]
        shift_value = override_shift if override_shift >= 0.0 else default_shift
        print(f"Selected Shift Value: {shift_value} (Override: {override_shift}, Default: {default_shift})")
        
        if scheduler != "Default for model":
            print(f"Replacing default scheduler ({original_scheduler_class}) with: {scheduler}")
            
            # Create a completely fresh scheduler instance to avoid any parameter leakage
            if scheduler == "UniPC":
                new_scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000, 
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_karras_sigmas=True
                )
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Exponential":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000, 
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_exponential_sigmas=True
                )
                pipe.scheduler = new_scheduler
        else:
            # Ensure we're using the original scheduler as specified in the model config
            print(f"Using model's default scheduler type: {original_scheduler_class} with shift={shift_value}")
            pipe.scheduler = get_scheduler_instance(original_scheduler_class, shift_value)
                
        # --- Generation Setup ---
        is_nf4_current = config.get("is_nf4", False)
        num_inference_steps = override_steps if override_steps >= 0 else config["num_inference_steps"]
        guidance_scale = override_cfg if override_cfg >= 0.0 else config["guidance_scale"]

        # Create the progress bar
        pbar = comfy.utils.ProgressBar(num_inference_steps)
        
        try:
            inference_device = comfy.model_management.get_torch_device()
        except Exception:
            inference_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            
        print(f"Creating Generator on: {inference_device}")
        generator = torch.Generator(device=inference_device).manual_seed(seed)
        print(f"\n--- Starting Generation ---")
        print(f"Model: {model_type}{' (uncensored)' if use_uncensored_llm else ''}, Res: {height}x{width}, Steps: {num_inference_steps}, CFG: {guidance_scale}, Shift: {pipe.scheduler.config.shift}, Seed: {seed}")
        print(f"Sequence lengths - CLIP-L: {max_length_clip_l}, OpenCLIP: {max_length_openclip}, T5: {max_length_t5}, Llama: {max_length_llama}")
        
        # --- Run Inference ---
        pipeline_output = None
        try:
            if not is_nf4_current:
                print(f"Ensuring pipe on: {inference_device} (Offload NOT enabled)")
                pipe.to(inference_device)
            else:
                print(f"Skipping pipe.to({inference_device}) (CPU offload enabled).")
                
            print("Executing pipeline inference...")
            
            # Use specific prompts for each encoder, falling back to primary prompt if empty
            prompt_clip_l = clip_l_prompt.strip() if clip_l_prompt.strip() else primary_prompt
            prompt_openclip = openclip_prompt.strip() if openclip_prompt.strip() else primary_prompt
            prompt_t5 = t5_prompt.strip() if t5_prompt.strip() else primary_prompt
            prompt_llama = llama_prompt.strip() if llama_prompt.strip() else primary_prompt
            
            print(f"Using per-encoder prompts:")
            print(f"  CLIP-L ({max_length_clip_l} tokens): {prompt_clip_l[:50]}{'...' if len(prompt_clip_l) > 50 else ''}")
            print(f"  OpenCLIP ({max_length_openclip} tokens): {prompt_openclip[:50]}{'...' if len(prompt_openclip) > 50 else ''}")
            print(f"  T5 ({max_length_t5} tokens): {prompt_t5[:50]}{'...' if len(prompt_t5) > 50 else ''}")
            print(f"  Llama ({max_length_llama} tokens): {prompt_llama[:50]}{'...' if len(prompt_llama) > 50 else ''}")

                        
            # Replace truly blank inputs with minimal period
            if not prompt_clip_l.strip():
                prompt_clip_l = "."
                
            if not prompt_openclip.strip():
                prompt_openclip = "."
                
            if not prompt_t5.strip():
                prompt_t5 = "."
                
            # Custom system prompt for blank LLM prompts to try to prevent LLM output noise
            custom_system_prompt = llm_system_prompt
            if not prompt_llama.strip():
                prompt_llama = "."
                custom_system_prompt = "You will only output a single period as your output '.'\nDo not add any other acknowledgement or extra text or data."
            
            # Ensure batch size consistency for multiple images
            if num_images > 1:
                print(f"Preparing for batch generation with {num_images} images...")
                # Create a list to store outputs
                output_images_list = []
                for i in range(num_images):
                    print(f"Generating image {i+1}/{num_images}...")
                    # Generate one image at a time to avoid batch size issues
                    with torch.inference_mode():
                        single_output = pipe(
                            prompt=prompt_clip_l, 
                            prompt_2=prompt_openclip,
                            prompt_3=prompt_t5,
                            prompt_4=prompt_llama,
                            negative_prompt=negative_prompt.strip() if negative_prompt else None,
                            height=height,
                            width=width,
                            guidance_scale=guidance_scale,
                            num_inference_steps=num_inference_steps,
                            num_images_per_prompt=1,
                            generator=generator,
                            max_sequence_length_clip_l=max_length_clip_l,
                            max_sequence_length_openclip=max_length_openclip,
                            max_sequence_length_t5=max_length_t5,
                            max_sequence_length_llama=max_length_llama,
                            llm_system_prompt=custom_system_prompt,
                            clip_l_scale=clip_l_weight,
                            openclip_scale=openclip_weight,
                            t5_scale=t5_weight,
                            llama_scale=llama_weight,
                            callback_on_step_end_tensor_inputs=["latents"],
                        )
                    output_images_list.extend(single_output.images)
                    pbar.update_absolute((i + 1) * num_inference_steps // num_images)
            else:
                with torch.inference_mode():
                    pipeline_output = pipe(
                            prompt=prompt_clip_l, 
                            prompt_2=prompt_openclip,
                            prompt_3=prompt_t5,
                            prompt_4=prompt_llama,
                            negative_prompt=negative_prompt.strip() if negative_prompt else None,
                            height=height,
                            width=width,
                            guidance_scale=guidance_scale,
                            num_inference_steps=num_inference_steps,
                            num_images_per_prompt=num_images,
                            generator=generator,
                            max_sequence_length_clip_l=max_length_clip_l,
                            max_sequence_length_openclip=max_length_openclip,
                            max_sequence_length_t5=max_length_t5,
                            max_sequence_length_llama=max_length_llama,
                            llm_system_prompt=custom_system_prompt,
                            clip_l_scale=clip_l_weight,
                            openclip_scale=openclip_weight,
                            t5_scale=t5_weight,
                            llama_scale=llama_weight,
                            callback_on_step_end_tensor_inputs=["latents"],
                        )
                    output_images_list = pipeline_output.images

            print("Pipeline inference finished.")

        except Exception as e:
            print(f"!!! ERROR during execution: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, height, width, 3)),)
        finally:
            pbar.update_absolute(num_inference_steps) # Update pbar regardless
        print("--- Generation Complete ---")
        
        # Robust output handling
        if output_images_list is None or not isinstance(output_images_list, list) or len(output_images_list) == 0:
            print(f"ERROR: No images returned or invalid format (Type: {type(output_images_list)}). Creating blank image.")
            return (torch.zeros((1, height, width, 3)),)

        try:
            print(f"Processing {len(output_images_list)} output image(s).")
            tensor_list = []

            for i, img in enumerate(output_images_list):
                if not isinstance(img, Image.Image):
                    print(f"WARNING: Item {i} in output list is not a PIL Image (Type: {type(img)}). Skipping.")
                    continue

                print(f"Converting image {i+1}/{len(output_images_list)}...")
                single_tensor = pil2tensor(img) # This returns shape [1, H, W, C]

                if single_tensor is not None:
                    if len(single_tensor.shape) == 4 and single_tensor.shape[0] == 1:
                        tensor_list.append(single_tensor)
                    else:
                        print(f"WARNING: pil2tensor returned unexpected shape {single_tensor.shape} for image {i}. Skipping.")
                else:
                    print(f"WARNING: pil2tensor failed for image {i}. Skipping.")

            if not tensor_list:
                print("ERROR: All image conversions failed. Creating blank image.")
                return (torch.zeros((1, height, width, 3)),)

            output_tensor = torch.cat(tensor_list, dim=0)
            print(f"Successfully converted {output_tensor.shape[0]} images into batch tensor.")

            if output_tensor.dtype != torch.float32:
                print(f"Converting batched {output_tensor.dtype} tensor to float32 for ComfyUI compatibility")
                output_tensor = output_tensor.to(torch.float32)

            if len(output_tensor.shape) != 4 or output_tensor.shape[0] == 0 or output_tensor.shape[3] != 3:
                print(f"ERROR: Invalid final batch tensor shape {output_tensor.shape}. Creating blank image.")
                return (torch.zeros((1, height, width, 3)),)

            print(f"Output tensor shape: {output_tensor.shape}") 

            try:
                import comfy.model_management as model_management
                print("HiDream: Requesting ComfyUI memory cleanup...")
                model_management.soft_empty_cache()
            except Exception as e:
                print(f"HiDream: ComfyUI cleanup failed: {e}")

            if torch.cuda.is_available():
                final_mem = torch.cuda.memory_allocated() / 1024**2
                print(f"HiDream: Final VRAM usage: {final_mem:.2f} MB (Change: {final_mem-initial_mem:.2f} MB)")
                
            return (output_tensor,)
            
        except Exception as e:
            print(f"Error processing output image: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, height, width, 3)),)


class HiDreamImg2Img:
    _model_cache = HiDreamSampler._model_cache
    cleanup_models = HiDreamSampler.cleanup_models
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate"
    CATEGORY = "HiDream"
    
    @classmethod
    def INPUT_TYPES(s):
        available_model_types = list(MODEL_CONFIGS.keys())
        if not available_model_types:
            return {"required": {"error": ("STRING", {"default": "No models available...", "multiline": True})}}
        
        default_model = "fast-nf4" if "fast-nf4" in available_model_types else "fast" if "fast" in available_model_types else available_model_types[0]
        
        # Define schedulers
        scheduler_options = [
            "Default for model",
            "UniPC",
            "Euler",
            "Karras Euler",
            "Karras Exponential"
        ]
        
        return {
            "required": {
                "model_type": (available_model_types, {"default": default_model}),
                "image": ("IMAGE",),
                "denoising_strength": ("FLOAT", {"default": 0.75, "min": 0.0, "max": 1.0, "step": 0.01}),
                "prompt": ("STRING", {"multiline": True, "default": "..."}),
                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "scheduler": (scheduler_options, {"default": "Default for model"}),
                "override_steps": ("INT", {"default": -1, "min": -1, "max": 100}),
                "override_cfg": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}),
                "override_shift": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 20.0, "step": 0.1}), # <-- Added Shift Override
                "use_uncensored_llm": ("BOOLEAN", {"default": False})
            },
            "optional": {
                "llm_system_prompt": ("STRING", {
                    "multiline": True,
                    "default": "You are a creative AI assistant that helps create detailed, vivid images based on user descriptions."
                }),
                "clip_l_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "openclip_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "t5_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
                "llama_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
            }
        }

    def preprocess_image(self, image, target_height=None, target_width=None):
        """Resize and possibly crop input image to match model requirements."""
        import torch.nn.functional as F
        import math
        
        # Get original dimensions
        _, orig_h, orig_w, _ = image.shape
        orig_aspect = orig_w / orig_h
        
        print(f"Original image dimensions: {orig_w}x{orig_h}, aspect ratio: {orig_aspect:.3f}")
        
        # If no target size provided, find closest standard resolution
        if target_height is None or target_width is None:
            # Define standard resolutions (must be divisible by 16)
            standard_resolutions = [
                (1024, 1024),  # 1:1
                (768, 1360),   # 9:16 (portrait)
                (1360, 768),   # 16:9 (landscape)
                (880, 1168),   # 3:4 (portrait)
                (1168, 880),   # 4:3 (landscape)
                (832, 1248),   # 2:3 (portrait)
                (1248, 832),   # 3:2 (landscape)
            ]
            
            # Find closest aspect ratio
            best_diff = float('inf')
            target_width, target_height = standard_resolutions[0]  # Default to square
            
            for w, h in standard_resolutions:
                res_aspect = w / h
                diff = abs(res_aspect - orig_aspect)
                if diff < best_diff:
                    best_diff = diff
                    target_width, target_height = w, h
            
            print(f"Selected target resolution: {target_width}x{target_height}")
        
        # Ensure dimensions are divisible by 16
        target_width = (target_width // 16) * 16
        target_height = (target_height // 16) * 16
        
        # Convert to format expected by F.interpolate [B,C,H,W]
        # ComfyUI typically uses [B,H,W,C]
        x = image.permute(0, 3, 1, 2)
        
        # Calculate resize dimensions preserving aspect ratio
        if orig_aspect > target_width / target_height:  # Image is wider
            new_w = target_width
            new_h = int(new_w / orig_aspect)
            new_h = (new_h // 16) * 16  # Make divisible by 16
        else:  # Image is taller
            new_h = target_height
            new_w = int(new_h * orig_aspect)
            new_w = (new_w // 16) * 16  # Make divisible by 16
        
        # Resize to preserve aspect ratio
        x_resized = F.interpolate(x, size=(new_h, new_w), mode='bicubic', align_corners=False)
        
        # Create target tensor with correct dimensions
        x_result = torch.zeros(1, 3, target_height, target_width, device=x.device, dtype=x.dtype)
        
        # Calculate position for center crop
        y_offset = max(0, (new_h - target_height) // 2)
        x_offset = max(0, (new_w - target_width) // 2)
        
        # Calculate how much to copy
        height_to_copy = min(new_h, target_height)
        width_to_copy = min(new_w, target_width)
        
        # Place the resized image in the center of the target tensor
        target_y_offset = max(0, (target_height - height_to_copy) // 2)
        target_x_offset = max(0, (target_width - width_to_copy) // 2)
        
        x_result[:, :, 
                 target_y_offset:target_y_offset+height_to_copy, 
                 target_x_offset:target_x_offset+width_to_copy] = x_resized[:, :, 
                                                                           y_offset:y_offset+height_to_copy, 
                                                                           x_offset:x_offset+width_to_copy]
        
        print(f"Processed to: {target_width}x{target_height} (divisible by 16)")
        
        # Convert back to ComfyUI format [B,H,W,C]
        return x_result.permute(0, 2, 3, 1)
    
    def generate(self, model_type, image, denoising_strength, prompt, negative_prompt, 
             seed, scheduler, override_steps, override_cfg, override_shift, use_uncensored_llm=False, # <-- Added override_shift
             llm_system_prompt="You are a creative AI assistant...",
             clip_l_weight=1.0, openclip_weight=1.0, t5_weight=1.0, llama_weight=1.0, **kwargs):


        # Preprocess the input image to ensure compatible dimensions
        processed_image = self.preprocess_image(image)
        
        # Get dimensions from processed image for the output
        _, height, width, _ = processed_image.shape
                 
        # Monitor initial memory usage
        if torch.cuda.is_available():
            initial_mem = torch.cuda.memory_allocated() / 1024**2
            print(f"HiDream: Initial VRAM usage: {initial_mem:.2f} MB")
            
        if not MODEL_CONFIGS or model_type == "error":
            print("HiDream Error: No models loaded.")
            return (torch.zeros((1, 512, 512, 3)),)
            
        pipe = None
        config = None
        
        # Create cache key that includes uncensored state
        cache_key = f"{model_type}_img2img_{'uncensored' if use_uncensored_llm else 'standard'}"
        
        # Try to reuse from cache first
        if cache_key in self._model_cache:
            print(f"Checking cache for {cache_key}...")
            pipe, config = self._model_cache[cache_key]
            valid_cache = True
            if pipe is None or config is None or not hasattr(pipe, 'transformer') or pipe.transformer is None:
                valid_cache = False
                print("Invalid cache, reloading...")
                del self._model_cache[cache_key]
                pipe, config = None, None
            if valid_cache:
                print("Using cached model.")
        
        if pipe is None:
            if self._model_cache:
                print(f"Clearing img2img cache before loading {model_type}...")
                keys_to_del = list(self._model_cache.keys())
                for key in keys_to_del:
                    print(f"  Removing '{key}'...")
                    try:
                        pipe_to_del, _= self._model_cache.pop(key)
                        # More aggressive cleanup
                        if hasattr(pipe_to_del, 'transformer'):
                            pipe_to_del.transformer = None
                        if hasattr(pipe_to_del, 'text_encoder_4'):
                            pipe_to_del.text_encoder_4 = None
                        if hasattr(pipe_to_del, 'tokenizer_4'):
                            pipe_to_del.tokenizer_4 = None
                        if hasattr(pipe_to_del, 'scheduler'):
                            pipe_to_del.scheduler = None
                        del pipe_to_del
                    except Exception as e:
                        print(f"  Error removing {key}: {e}")
                
                # Multiple garbage collection passes
                for _ in range(3):
                    gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    # Force synchronization
                    torch.cuda.synchronize()
                print("Cache cleared.")
            
            print(f"Loading model for {model_type} img2img...")

            try:
                # First load regular model
                txt2img_pipe, config = load_models(model_type, use_uncensored_llm)
                
                # Convert to img2img pipeline
                print("Creating img2img pipeline from loaded txt2img pipeline...")
                pipe = HiDreamImageToImagePipeline(
                    scheduler=txt2img_pipe.scheduler,
                    vae=txt2img_pipe.vae,
                    text_encoder=txt2img_pipe.text_encoder,
                    tokenizer=txt2img_pipe.tokenizer,
                    text_encoder_2=txt2img_pipe.text_encoder_2,
                    tokenizer_2=txt2img_pipe.tokenizer_2,
                    text_encoder_3=txt2img_pipe.text_encoder_3,
                    tokenizer_3=txt2img_pipe.tokenizer_3,
                    text_encoder_4=txt2img_pipe.text_encoder_4,
                    tokenizer_4=txt2img_pipe.tokenizer_4,
                )
                
                # Copy transformer and move to right device
                pipe.transformer = txt2img_pipe.transformer
                
                # Cleanup txt2img pipeline references
                txt2img_pipe = None
                
                # Cache the img2img pipeline
                self._model_cache[cache_key] = (pipe, config)
                print(f"Model {model_type} loaded & cached for img2img!")
                
            except Exception as e:
                print(f"!!! ERROR loading {model_type}: {e}")
                import traceback
                traceback.print_exc()
                return (torch.zeros((1, 512, 512, 3)),)
                
        if pipe is None or config is None:
            print("CRITICAL ERROR: Load failed.")
            return (torch.zeros((1, 512, 512, 3)),)
        
        # Update scheduler if requested
        original_scheduler_class = config["scheduler_class"]
        # Determine shift value: override if >= 0.0, else use config default
        default_shift = config["shift"]
        shift_value = override_shift if override_shift >= 0.0 else default_shift
        print(f"Selected Shift Value: {shift_value} (Override: {override_shift}, Default: {default_shift})")
             
        if scheduler != "Default for model":
            print(f"Replacing default scheduler ({original_scheduler_class}) with: {scheduler}")
            # Create a completely fresh scheduler instance to avoid any parameter leakage
            if scheduler == "UniPC":
                new_scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift_value, use_dynamic_shifting=False)
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Euler":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000,
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_karras_sigmas=True
                )
                pipe.scheduler = new_scheduler
            elif scheduler == "Karras Exponential":
                new_scheduler = FlashFlowMatchEulerDiscreteScheduler(
                    num_train_timesteps=1000,
                    shift=shift_value,
                    use_dynamic_shifting=False,
                    use_exponential_sigmas=True
                )
                pipe.scheduler = new_scheduler
            else: # Should not happen
                 print(f"Warning: Unknown scheduler option '{scheduler}'. Using default.")
                 pipe.scheduler = get_scheduler_instance(original_scheduler_class, shift_value)
                
        else:
            # Ensure we're using the original scheduler as specified in the model config
            print(f"Using model's default scheduler: {original_scheduler_class} with shift={shift_value}")
            pipe.scheduler = get_scheduler_instance(original_scheduler_class, shift_value)
            
        # Setup generation parameters
        is_nf4_current = config.get("is_nf4", False) 
        num_inference_steps = override_steps if override_steps >= 0 else config["num_inference_steps"]
        guidance_scale = override_cfg if override_cfg >= 0.0 else config["guidance_scale"]
        
        # Create progress bar
        pbar = comfy.utils.ProgressBar(num_inference_steps)
        
        # Define progress callback
        def progress_callback(pipe, i, t, callback_kwargs):
            # Update ComfyUI progress bar
            pbar.update_absolute(i+1)
            return callback_kwargs
            
        try:
            inference_device = comfy.model_management.get_torch_device()
        except Exception:
            inference_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            
        print(f"Creating Generator on: {inference_device}")
        generator = torch.Generator(device=inference_device).manual_seed(seed)
        print(f"\n--- Starting Img2Img Generation ---")
        _, h, w, _ = image.shape
        print(f"Model: {model_type}{' (uncensored)' if use_uncensored_llm else ''}, Input Size: {h}x{w}")
        print(f"Denoising: {denoising_strength}, Steps: {num_inference_steps}, CFG: {guidance_scale}, Shift: {pipe.scheduler.config.shift}, Seed: {seed}")

        # --- Run Inference ---
        output_images = None
        try:
            if not is_nf4_current:
                print(f"Ensuring pipe on: {inference_device} (Offload NOT enabled)")
                pipe.to(inference_device)
            else:
                print(f"Skipping pipe.to({inference_device}) (CPU offload enabled).")
                
            print("Executing pipeline inference...")
            
            with torch.inference_mode():
                output_images = pipe(
                    prompt=prompt,
                    prompt_2=prompt,  # Same prompt for all encoders
                    prompt_3=prompt,
                    prompt_4=prompt,
                    negative_prompt=negative_prompt.strip() if negative_prompt else None,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    num_images_per_prompt=1,
                    generator=generator,
                    init_image=processed_image,
                    denoising_strength=denoising_strength,
                    llm_system_prompt=llm_system_prompt,
                    clip_l_scale=clip_l_weight,
                    openclip_scale=openclip_weight,
                    t5_scale=t5_weight,
                    llama_scale=llama_weight,
                    callback_on_step_end=progress_callback,
                    callback_on_step_end_tensor_inputs=["latents"],
                ).images
                
            print("Pipeline inference finished.")
            
        except Exception as e:
            print(f"!!! ERROR during execution: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, h, w, 3)),)
        finally:
            pbar.update_absolute(num_inference_steps) # Update pbar regardless
        print("--- Generation Complete ---")
        
        # Robust output handling
        if output_images is None or len(output_images) == 0:
            print("ERROR: No images returned. Creating blank image.")
            return (torch.zeros((1, h, w, 3)),)
            
        try:
            print(f"Processing output image. Type: {type(output_images[0])}")
            output_tensor = pil2tensor(output_images[0])
            
            if output_tensor is None:
                print("ERROR: pil2tensor returned None. Creating blank image.")
                return (torch.zeros((1, h, w, 3)),)
                
            # Fix for bfloat16 tensor issue
            if output_tensor.dtype == torch.bfloat16:
                print("Converting bfloat16 tensor to float32 for ComfyUI compatibility")
                output_tensor = output_tensor.to(torch.float32)
                
            # Verify tensor shape is valid
            if len(output_tensor.shape) != 4 or output_tensor.shape[0] != 1 or output_tensor.shape[3] != 3:
                print(f"ERROR: Invalid tensor shape {output_tensor.shape}. Creating blank image.")
                return (torch.zeros((1, h, w, 3)),)
                
            print(f"Output tensor shape: {output_tensor.shape}")
            
            # After generating the image, try to clean up any temporary memory
            try:
                import comfy.model_management as model_management
                print("HiDream: Requesting ComfyUI memory cleanup...")
                model_management.soft_empty_cache()
            except Exception as e:
                print(f"HiDream: ComfyUI cleanup failed: {e}")
                
            # Log final memory usage
            if torch.cuda.is_available():
                final_mem = torch.cuda.memory_allocated() / 1024**2
                print(f"HiDream: Final VRAM usage: {final_mem:.2f} MB (Change: {final_mem-initial_mem:.2f} MB)")
                
            return (output_tensor,)
            
        except Exception as e:
            print(f"Error processing output image: {e}")
            import traceback
            traceback.print_exc()
            return (torch.zeros((1, h, w, 3)),)

# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
    "HiDreamSampler": HiDreamSampler,
    "HiDreamSamplerAdvanced": HiDreamSamplerAdvanced,
    "HiDreamImg2Img": HiDreamImg2Img
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "HiDreamSampler": "HiDream Sampler",
    "HiDreamSamplerAdvanced": "HiDream Sampler (Advanced)",
    "HiDreamImg2Img": "HiDream Image to Image"
}

# --- Register with ComfyUI's Memory Management ---
try:
    import comfy.model_management as model_management
    
    # Check if we can register a cleanup callback
    if hasattr(model_management, 'unload_all_models'):
        original_unload = model_management.unload_all_models
        
        # Wrap the original function to include our cleanup
        def wrapped_unload():
            print("HiDream: ComfyUI is unloading all models, cleaning HiDream cache...")
            HiDreamSampler.cleanup_models()
            return original_unload()
            
        # Replace the original function with our wrapped version
        model_management.unload_all_models = wrapped_unload
        print("HiDream: Successfully registered with ComfyUI memory management")
except Exception as e:
    print(f"HiDream: Could not register cleanup with model_management: {e}")

print("-" * 50 + "\nHiDream Sampler Node Initialized\nAvailable Models: " + str(list(MODEL_CONFIGS.keys())) + "\n" + "-" * 50)