# Copyright (c) 2026 exportAnything. All rights reserved. # SPDX-License-Identifier: MIT from __future__ import annotations import argparse import importlib.util import json import sys from pathlib import Path COMFY_ROOT = Path(__file__).resolve().parents[2] if str(COMFY_ROOT) not in sys.path: sys.path.append(str(COMFY_ROOT)) def _read_json_file(path: Path) -> dict: try: data = json.loads(path.read_text(encoding="utf-8")) return data if isinstance(data, dict) else {} except Exception: return {} def _model_path_info(model_path: str) -> dict: path = Path(model_path or "") info = { "path_kind": "missing", "is_gguf_file": False, "is_hf_repo": False, "model_type": "", "architecture": "", "quant_algo": "", "quant_method": "", "quant_producer": "", "quant_producer_version": "", "transformers_version_required": "", "has_generation_config": False, "has_modelopt_state": False, "runtime_hint": "", } if not path.exists(): return info if path.is_file(): info["path_kind"] = "gguf_file" if path.suffix.lower() == ".gguf" else "file" info["is_gguf_file"] = path.suffix.lower() == ".gguf" if info["is_gguf_file"]: info["runtime_hint"] = "llama-diffusion-cli GGUF" return info if not path.is_dir(): info["path_kind"] = "unknown" return info config = _read_json_file(path / "config.json") hf_quant_config = _read_json_file(path / "hf_quant_config.json") quant_config = config.get("quantization_config") if isinstance(config.get("quantization_config"), dict) else {} if not quant_config: quant_config = hf_quant_config.get("quantization") if isinstance(hf_quant_config.get("quantization"), dict) else {} producer = quant_config.get("producer") if isinstance(quant_config.get("producer"), dict) else {} if not producer: producer = hf_quant_config.get("producer") if isinstance(hf_quant_config.get("producer"), dict) else {} architectures = config.get("architectures") if isinstance(config.get("architectures"), list) else [] info.update( { "is_hf_repo": bool(config), "model_type": str(config.get("model_type", "")), "architecture": str(architectures[0]) if architectures else "", "quant_algo": str(quant_config.get("quant_algo", "")), "quant_method": str(quant_config.get("quant_method", "")), "quant_producer": str(producer.get("name", "")), "quant_producer_version": str(producer.get("version", "")), "transformers_version_required": str(config.get("transformers_version", "")), "has_generation_config": (path / "generation_config.json").exists(), "has_modelopt_state": (path / "modelopt_state.pth").exists(), } ) if info["quant_algo"].upper() == "NVFP4": info["path_kind"] = "nvfp4_hf_repo" info["runtime_hint"] = "NVIDIA ModelOpt NVFP4 checkpoint" elif config: info["path_kind"] = "hf_repo" info["runtime_hint"] = "Transformers whole-repo model" else: info["path_kind"] = "directory" return info def _short_video_processor_smoke(model_path: str) -> dict: result = { "passed": False, "frame_count": 11, "transport": "", "input_keys": [], "pixel_values_shape": [], "pixel_values_videos_shape": [], "error": "", } try: import torch from transformers import AutoProcessor from nodes import MediaContext, _transformers_messages_and_processor_kwargs processor = AutoProcessor.from_pretrained(model_path, local_files_only=True) frames = torch.zeros((11, 64, 64, 3), dtype=torch.float32) frames[:, :, :, 0] = torch.linspace(0.0, 1.0, 11).view(11, 1, 1) context = MediaContext( images=frames, source="video", metadata={ "source": "video", "duration_seconds": 11.0, "sample_fps": 1.0, "sampled_frame_count": 11, "width": 64, "height": 64, }, ) messages, processor_kwargs = _transformers_messages_and_processor_kwargs("Describe this video.", context) inputs = processor.apply_chat_template( messages, tokenize=True, return_dict=True, return_tensors="pt", add_generation_prompt=True, processor_kwargs=processor_kwargs, ) result["input_keys"] = sorted(str(key) for key in inputs.keys()) image_values = inputs.get("pixel_values") video_values = inputs.get("pixel_values_videos") result["pixel_values_shape"] = list(image_values.shape) if image_values is not None else [] result["pixel_values_videos_shape"] = list(video_values.shape) if video_values is not None else [] if image_values is not None: result["transport"] = "sampled_frame_images" result["passed"] = bool(image_values.shape[0] == 11) elif video_values is not None: result["transport"] = "video_tokens" result["passed"] = bool(video_values.shape[1] == 11) except Exception as exc: result["error"] = str(exc) return result def probe(model_path: str) -> dict: result = { "python": sys.executable, "model_path": model_path, "model_path_exists": Path(model_path).exists(), **_model_path_info(model_path), "transformers_importable": importlib.util.find_spec("transformers") is not None, "diffusion_gemma_module": importlib.util.find_spec("transformers.models.diffusion_gemma") is not None, "vllm_importable": importlib.util.find_spec("vllm") is not None, "modelopt_importable": bool( importlib.util.find_spec("modelopt") or importlib.util.find_spec("nvidia_modelopt") ), "comfy_kitchen_importable": importlib.util.find_spec("comfy_kitchen") is not None, "fp_quant_importable": importlib.util.find_spec("fp_quant") is not None, "qutlass_importable": importlib.util.find_spec("qutlass") is not None, "tensorrt_llm_importable": importlib.util.find_spec("tensorrt_llm") is not None, "torch_cuda_available": False, "torch_cuda_capability": "", "transformers_version": "", "processor_class": "", "auto_processor": False, "diffusion_gemma_class": False, "processor_loads": False, "modelopt_nvfp4_supported": False, "comfy_nvfp4_bridge_supported": False, "nvfp4_bridge_smoke": {}, "short_video_processor_smoke": {}, "notes": [], "errors": [], } try: import torch result["torch_cuda_available"] = bool(torch.cuda.is_available()) if torch.cuda.is_available(): result["torch_cuda_capability"] = ".".join(str(part) for part in torch.cuda.get_device_capability()) except Exception as exc: result["errors"].append(f"Torch CUDA probe failed: {exc}") try: import transformers from transformers import AutoProcessor result["transformers_version"] = str(getattr(transformers, "__version__", "")) result["auto_processor"] = True if result["model_path_exists"]: processor = AutoProcessor.from_pretrained(model_path, local_files_only=True) result["processor_class"] = type(processor).__name__ result["processor_loads"] = True result["short_video_processor_smoke"] = _short_video_processor_smoke(model_path) if not result["short_video_processor_smoke"].get("passed"): result["errors"].append( "Short sampled-video processor smoke failed; video pixels may fall back before model generation." ) except Exception as exc: result["errors"].append(f"AutoProcessor probe failed: {exc}") try: from transformers import DiffusionGemmaForBlockDiffusion # noqa: F401 result["diffusion_gemma_class"] = True except Exception as exc: result["errors"].append(f"DiffusionGemmaForBlockDiffusion probe failed: {exc}") if result["path_kind"] == "nvfp4_hf_repo": modelopt_nvfp4 = result["quant_method"] == "modelopt" try: from nodes import RuntimeConfig, _runtime_status runtime_status = _runtime_status( RuntimeConfig( model_path=model_path, backend="transformers_inprocess", dtype="auto", quantization="modelopt_nvfp4", local_files_only=True, unload_policy="keep_loaded", max_memory_gb=30.0, ) ) result["comfy_nvfp4_bridge_supported"] = bool(runtime_status.get("comfy_nvfp4_bridge_supported")) result["nvfp4_bridge_smoke"] = runtime_status.get("nvfp4_bridge_smoke", {}) result["modelopt_nvfp4_supported"] = bool(runtime_status.get("modelopt_nvfp4_supported")) except Exception as exc: result["errors"].append(f"Comfy NVFP4 bridge probe failed: {exc}") if result["comfy_nvfp4_bridge_supported"]: result["notes"].append( "NVFP4 Hugging Face repo detected. Transformers can parse DiffusionGemma, and the local " "Comfy NVFP4 bridge smoke passed for ModelOpt packed expert weights." ) else: result["modelopt_nvfp4_supported"] = bool( modelopt_nvfp4 and result["modelopt_importable"] and result["tensorrt_llm_importable"] and result["has_modelopt_state"] ) result["notes"].append( "NVFP4 Hugging Face repo detected. Transformers can parse the DiffusionGemma config and processor, " "but this ModelOpt NVFP4 checkpoint still requires a supported in-process ModelOpt NVFP4 loader." ) if modelopt_nvfp4 and not result["modelopt_nvfp4_supported"]: result["errors"].append( "ModelOpt NVFP4 load is blocked: Transformers does not register quant_method=modelopt, " "modelopt_state.pth is absent, TensorRT-LLM is unavailable, or the Comfy NVFP4 bridge did not pass." ) result["gate_1_passed"] = bool( result["model_path_exists"] and result["auto_processor"] and result["diffusion_gemma_class"] and result["processor_loads"] and not ( result["path_kind"] == "nvfp4_hf_repo" and result["quant_method"] == "modelopt" and not result["modelopt_nvfp4_supported"] ) ) result["class_gate_passed"] = bool( result["model_path_exists"] and result["auto_processor"] and result["diffusion_gemma_class"] and result["processor_loads"] ) return result def main() -> int: parser = argparse.ArgumentParser(description="DiffusionGemma runtime proof-gate probe.") parser.add_argument( "--model-path", default=str(Path.cwd() / "models" / "LLM" / "diffusiongemma-26B-A4B-it"), help="Local whole-repo DiffusionGemma path.", ) args = parser.parse_args() print(json.dumps(probe(args.model_path), indent=2, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())