checkpoint_multigpu.py

import torch
import logging
import hashlib
import comfy.sd
import comfy.utils
import comfy.model_management as mm
import comfy.model_detection
import comfy.clip_vision
from comfy.sd import VAE, CLIP
from .device_utils import get_device_list, soft_empty_cache_multigpu
from .model_management_mgpu import multigpu_memory_log
from .distorch_2 import register_patched_safetensor_modelpatcher

logger = logging.getLogger("MultiGPU")

checkpoint_device_config = {}
checkpoint_distorch_config = {}

original_load_state_dict_guess_config = None

def patch_load_state_dict_guess_config():
    """Monkey patch comfy.sd.load_state_dict_guess_config with MultiGPU-aware checkpoint loading."""
    global original_load_state_dict_guess_config

    if original_load_state_dict_guess_config is not None:
        logger.debug("[MultiGPU Checkpoint] load_state_dict_guess_config is already patched.")
        return

    logger.info("[MultiGPU Core Patching] Patching comfy.sd.load_state_dict_guess_config for advanced MultiGPU loading.")
    original_load_state_dict_guess_config = comfy.sd.load_state_dict_guess_config
    comfy.sd.load_state_dict_guess_config = patched_load_state_dict_guess_config

def patched_load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False,
                                        embedding_directory=None, output_model=True, model_options={},
                                        te_model_options={}, metadata=None, disable_dynamic=False):
    """Patched checkpoint loader with MultiGPU and DisTorch2 device placement support."""
    from . import set_current_device, set_current_text_encoder_device, get_current_device, get_current_text_encoder_device

    sd_size = sum(p.numel() for p in sd.values() if hasattr(p, 'numel'))
    config_hash = str(sd_size)
    device_config = checkpoint_device_config.get(config_hash)
    distorch_config = checkpoint_distorch_config.get(config_hash)

    if not device_config and not distorch_config:
        return original_load_state_dict_guess_config(
            sd,
            output_vae=output_vae,
            output_clip=output_clip,
            output_clipvision=output_clipvision,
            embedding_directory=embedding_directory,
            output_model=output_model,
            model_options=model_options,
            te_model_options=te_model_options,
            metadata=metadata,
            disable_dynamic=disable_dynamic,
        )

    logger.debug("[MultiGPU Checkpoint] ENTERING Patched Checkpoint Loader")
    logger.debug(f"[MultiGPU Checkpoint] Received Device Config: {device_config}")
    logger.debug(f"[MultiGPU Checkpoint] Received DisTorch2 Config: {distorch_config}")

    clip = None
    clipvision = None
    vae = None
    model = None
    model_patcher = None

    # Capture the current devices at runtime so we can restore them after loading
    original_main_device = get_current_device()
    original_clip_device = get_current_text_encoder_device()

    try:
        diffusion_model_prefix = comfy.model_detection.unet_prefix_from_state_dict(sd)
        parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
        weight_dtype = comfy.utils.weight_dtype(sd, diffusion_model_prefix)

        custom_operations = model_options.get("custom_operations", None)
        if custom_operations is None:
            sd, metadata = comfy.utils.convert_old_quants(sd, diffusion_model_prefix, metadata=metadata)

        model_config = comfy.model_detection.model_config_from_unet(sd, diffusion_model_prefix, metadata=metadata)

        if model_config is None:
            logger.warning("[MultiGPU] Warning: Not a standard checkpoint file. Trying to load as diffusion model only.")
            # Simplified fallback for non-checkpoints
            set_current_device(device_config.get('unet_device', original_main_device))
            diffusion_model = comfy.sd.load_diffusion_model_state_dict(
                sd,
                model_options={},
                metadata=metadata,
                disable_dynamic=disable_dynamic,
            )
            if diffusion_model is None:
                return None
            return (diffusion_model, None, VAE(sd={}), None)

        logger.debug(f"[MultiGPU] Detected Model Config: {type(model_config).__name__}, Parameters: {parameters/10**9:.2f}B")

        unet_weight_dtype = list(model_config.supported_inference_dtypes)
        if model_config.scaled_fp8 is not None:
            weight_dtype = None

        if custom_operations is not None:
            model_config.custom_operations = custom_operations
        unet_dtype = model_options.get("dtype", model_options.get("weight_dtype", None))
        if unet_dtype is None:
            unet_dtype = mm.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)

        unet_compute_device = torch.device(device_config.get('unet_device', original_main_device))
        if model_config.scaled_fp8 is not None:
            manual_cast_dtype = mm.unet_manual_cast(None, unet_compute_device, model_config.supported_inference_dtypes)
        else:
            manual_cast_dtype = mm.unet_manual_cast(unet_dtype, unet_compute_device, model_config.supported_inference_dtypes)
        model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
        logger.info(f"UNet DType: {unet_dtype}, Manual Cast: {manual_cast_dtype}")


        if model_config.clip_vision_prefix is not None and output_clipvision:
            clipvision = comfy.clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)

        if output_model:
            unet_compute_device = torch.device(device_config.get('unet_device', original_main_device))
            set_current_device(unet_compute_device)
            inital_load_device = mm.unet_inital_load_device(parameters, unet_dtype)

            multigpu_memory_log(f"unet:{config_hash[:8]}", "pre-load")

            model = model_config.get_model(sd, diffusion_model_prefix, device=inital_load_device)
            model_patcher_class = comfy.model_patcher.ModelPatcher if disable_dynamic else comfy.model_patcher.CoreModelPatcher
            model_patcher = model_patcher_class(model, load_device=unet_compute_device, offload_device=mm.unet_offload_device())
            model.load_model_weights(sd, diffusion_model_prefix, assign=model_patcher.is_dynamic())
            multigpu_memory_log(f"unet:{config_hash[:8]}", "post-weights")

            logger.mgpu_mm_log("Invoking soft_empty_cache_multigpu before UNet ModelPatcher setup")
            soft_empty_cache_multigpu()
            multigpu_memory_log(f"unet:{config_hash[:8]}", "post-model")

            if distorch_config and 'unet_allocation' in distorch_config:
                unet_alloc = distorch_config['unet_allocation']
                if unet_alloc:
                    register_patched_safetensor_modelpatcher()
                    inner_model = model_patcher.model
                    inner_model._distorch_v2_meta = {"full_allocation": unet_alloc}
                    logger.info(f"[CHECKPOINT_META] UNET inner_model id=0x{id(inner_model):x}")
                    model._distorch_high_precision_loras = distorch_config.get('high_precision_loras', True)

        if output_vae:
            vae_target_device = torch.device(device_config.get('vae_device', original_main_device))
            set_current_device(vae_target_device) # Use main device context for VAE
            multigpu_memory_log(f"vae:{config_hash[:8]}", "pre-load")

            vae_sd = comfy.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)
            vae_sd = model_config.process_vae_state_dict(vae_sd)
            vae = VAE(sd=vae_sd, metadata=metadata)
            multigpu_memory_log(f"vae:{config_hash[:8]}", "post-load")

        if output_clip:
            if te_model_options.get("custom_operations", None) is None:
                scaled_fp8_list = []
                for k in list(sd.keys()):  # Convert scaled fp8 to mixed ops
                    if k.endswith(".scaled_fp8"):
                        scaled_fp8_list.append(k[:-len("scaled_fp8")])

                if len(scaled_fp8_list) > 0:
                    out_sd = {}
                    for k in sd:
                        skip = False
                        for pref in scaled_fp8_list:
                            skip = skip or k.startswith(pref)
                        if not skip:
                            out_sd[k] = sd[k]

                    for pref in scaled_fp8_list:
                        quant_sd, qmetadata = comfy.utils.convert_old_quants(sd, pref, metadata={})
                        for k in quant_sd:
                            out_sd[k] = quant_sd[k]
                        sd = out_sd

            clip_target_device = torch.device(device_config.get('clip_device', original_clip_device))
            set_current_text_encoder_device(clip_target_device)

            clip_target = model_config.clip_target(state_dict=sd)
            if clip_target is not None:
                clip_sd = model_config.process_clip_state_dict(sd)
                if len(clip_sd) > 0:
                    logger.debug("[MultiGPU Checkpoint] Invoking soft_empty_cache_multigpu before CLIP construction")
                    multigpu_memory_log(f"clip:{config_hash[:8]}", "pre-load")
                    soft_empty_cache_multigpu()
                    clip_params = comfy.utils.calculate_parameters(clip_sd)
                    clip = CLIP(
                        clip_target,
                        embedding_directory=embedding_directory,
                        tokenizer_data=clip_sd,
                        parameters=clip_params,
                        state_dict=clip_sd,
                        model_options=te_model_options,
                        disable_dynamic=disable_dynamic,
                    )

                    if distorch_config and 'clip_allocation' in distorch_config:
                        clip_alloc = distorch_config['clip_allocation']
                        if clip_alloc and hasattr(clip, 'patcher'):
                            register_patched_safetensor_modelpatcher()
                            inner_clip = clip.patcher.model
                            inner_clip._distorch_v2_meta = {"full_allocation": clip_alloc}
                            logger.info(f"[CHECKPOINT_META] CLIP inner_model id=0x{id(inner_clip):x}")
                            clip.patcher.model._distorch_high_precision_loras = distorch_config.get('high_precision_loras', True)

                    logger.info("CLIP Loaded.")
                    multigpu_memory_log(f"clip:{config_hash[:8]}", "post-load")
                else:
                    logger.warning("No CLIP/text encoder weights in checkpoint.")
            else:
                logger.warning("CLIP target not found in model config.")

    finally:
        set_current_device(original_main_device)
        set_current_text_encoder_device(original_clip_device)
        if config_hash in checkpoint_device_config:
            del checkpoint_device_config[config_hash]
        if config_hash in checkpoint_distorch_config:
            del checkpoint_distorch_config[config_hash]
    return (model_patcher, clip, vae, clipvision)

class CheckpointLoaderAdvancedMultiGPU:
    @classmethod
    def INPUT_TYPES(s):
        import folder_paths
        devices = get_device_list()
        default_device = devices[1] if len(devices) > 1 else devices[0]

        return {
            "required": {
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
                "unet_device": (devices, {"default": default_device}),
                "clip_device": (devices, {"default": default_device}),
                "vae_device": (devices, {"default": default_device}),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"
    CATEGORY = "multigpu"
    TITLE = "Checkpoint Loader Advanced (MultiGPU)"

    def load_checkpoint(self, ckpt_name, unet_device, clip_device, vae_device):
        patch_load_state_dict_guess_config()

        import folder_paths
        import comfy.utils

        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
        sd = comfy.utils.load_torch_file(ckpt_path)
        sd_size = sum(p.numel() for p in sd.values() if hasattr(p, 'numel'))
        config_hash = str(sd_size)

        checkpoint_device_config[config_hash] = {
            'unet_device': unet_device, 'clip_device': clip_device, 'vae_device': vae_device
        }

        # Load using standard loader, our patch will intercept
        from nodes import CheckpointLoaderSimple
        return CheckpointLoaderSimple().load_checkpoint(ckpt_name)


class CheckpointLoaderAdvancedDisTorch2MultiGPU:
    @classmethod
    def INPUT_TYPES(s):
        import folder_paths
        devices = get_device_list()
        compute_device = devices[1] if len(devices) > 1 else devices[0]

        return {
            "required": {
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
                "unet_compute_device": (devices, {"default": compute_device}),
                "unet_virtual_vram_gb": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}),
                "unet_donor_device": (devices, {"default": "cpu"}),
                "clip_compute_device": (devices, {"default": "cpu"}),
                "clip_virtual_vram_gb": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 128.0, "step": 0.1}),
                "clip_donor_device": (devices, {"default": "cpu"}),
                "vae_device": (devices, {"default": compute_device}),
            }, "optional": {
                "unet_expert_mode_allocations": ("STRING", {"multiline": False, "default": ""}),
                "clip_expert_mode_allocations": ("STRING", {"multiline": False, "default": ""}),
                "high_precision_loras": ("BOOLEAN", {"default": True}),
                "eject_models": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"
    CATEGORY = "multigpu/distorch_2"
    TITLE = "Checkpoint Loader Advanced (DisTorch2)"

    def load_checkpoint(self, ckpt_name, unet_compute_device, unet_virtual_vram_gb, unet_donor_device,
                       clip_compute_device, clip_virtual_vram_gb, clip_donor_device, vae_device,
                       unet_expert_mode_allocations="", clip_expert_mode_allocations="", high_precision_loras=True, eject_models=True):

        if eject_models:
            logger.mgpu_mm_log("[EJECT_MODELS_SETUP] eject_models=True - marking all loaded models for eviction")
            ejection_count = 0
            for i, lm in enumerate(mm.current_loaded_models):
                model_name = type(getattr(lm.model, 'model', lm.model)).__name__ if lm.model else 'Unknown'
                if hasattr(lm.model, 'model') and lm.model.model is not None:
                    lm.model.model._mgpu_unload_distorch_model = True
                    logger.mgpu_mm_log(f"[EJECT_MARKED] Model {i}: {model_name} (id=0x{id(lm):x}) → marked for eviction")
                    ejection_count += 1
                elif lm.model is not None:
                    lm.model._mgpu_unload_distorch_model = True
                    logger.mgpu_mm_log(f"[EJECT_MARKED] Model {i}: {model_name} (direct patcher) → marked for eviction")
                    ejection_count += 1
            logger.mgpu_mm_log(f"[EJECT_MODELS_SETUP_COMPLETE] Marked {ejection_count} models for Comfy Core eviction during load_models_gpu")

        patch_load_state_dict_guess_config()

        import folder_paths
        import comfy.utils

        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
        sd = comfy.utils.load_torch_file(ckpt_path)
        sd_size = sum(p.numel() for p in sd.values() if hasattr(p, 'numel'))
        config_hash = str(sd_size)

        checkpoint_device_config[config_hash] = {
            'unet_device': unet_compute_device,
            'clip_device': clip_compute_device,
            'vae_device': vae_device
        }

        unet_vram_str = ""
        if unet_virtual_vram_gb > 0:
            unet_vram_str = f"{unet_compute_device};{unet_virtual_vram_gb};{unet_donor_device}"
        elif unet_expert_mode_allocations:
            unet_vram_str = unet_compute_device
        unet_alloc = f"{unet_expert_mode_allocations}#{unet_vram_str}" if unet_expert_mode_allocations or unet_vram_str else ""

        clip_vram_str = ""
        if clip_virtual_vram_gb > 0:
            clip_vram_str = f"{clip_compute_device};{clip_virtual_vram_gb};{clip_donor_device}"
        elif clip_expert_mode_allocations:
            clip_vram_str = clip_compute_device
        clip_alloc = f"{clip_expert_mode_allocations}#{clip_vram_str}" if clip_expert_mode_allocations or clip_vram_str else ""

        checkpoint_distorch_config[config_hash] = {
            'unet_allocation': unet_alloc,
            'clip_allocation': clip_alloc,
            'high_precision_loras': high_precision_loras,
            'unet_settings': hashlib.sha256(f"{unet_alloc}{high_precision_loras}".encode()).hexdigest(),
            'clip_settings': hashlib.sha256(f"{clip_alloc}{high_precision_loras}".encode()).hexdigest(),
        }

        from nodes import CheckpointLoaderSimple
        return CheckpointLoaderSimple().load_checkpoint(ckpt_name)