diff --git a/nodes/quantizer.py b/nodes/quantizer.py index 1d239fe..04dbc0a 100644 --- a/nodes/quantizer.py +++ b/nodes/quantizer.py @@ -134,11 +134,21 @@ class ModelOptQuantizeUNet: print(f"\nIntrospecting model architecture...") model_info = introspect_diffusion_model(diffusion_model) print(f" Architecture: {model_info['architecture']}") + print(f" Parameters: {model_info['param_count_billions']:.2f}B ({model_info['param_count']:,})") print(f" Y dimension: {model_info['y_dim']} ({model_info.get('y_dim_source', 'not detected')})") print(f" Context dimension: {model_info['context_dim']}") print(f" Latent format: {model_info['latent_channels']}x{model_info['latent_spatial']}x{model_info['latent_spatial']}") print(f" Has y parameter: {model_info['has_y_param']}") + # Debug: Print found attributes + if model_info.get('important_attributes_found'): + found = [k for k, v in model_info['important_attributes_found'].items() if v] + print(f" Found attributes: {', '.join(found) if found else 'None'}") + + # Debug: Print detection events + if model_info.get('detected_attributes'): + print(f" Detection events: {', '.join(model_info['detected_attributes'])}") + # Parse skip layers skip_layer_list = [] if skip_layers: @@ -172,6 +182,13 @@ class ModelOptQuantizeUNet: print(f" Using base config: {precision.upper()}") print(f" Algorithm: {quant_cfg.get('algorithm', 'default')}") + # Debug: Print config details to troubleshoot "Inserted 0 quantizers" + if "quant_cfg" in quant_cfg: + print(f" Config keys: {list(quant_cfg['quant_cfg'].keys())}") + # Check if there are any wildcard patterns + wildcards = [k for k in quant_cfg['quant_cfg'].keys() if '*' in k] + print(f" Wildcard patterns: {wildcards}") + # Customize config to skip certain layers if specified if skip_layer_list: import copy diff --git a/nodes/utils.py b/nodes/utils.py index c7dc876..d7681f5 100644 --- a/nodes/utils.py +++ b/nodes/utils.py @@ -297,8 +297,19 @@ def introspect_diffusion_model(model): "has_y_param": False, "model_channels": None, "num_heads": None, + "detected_attributes": [], # Track what we actually found } + # DEBUG: List all top-level attributes to help troubleshoot + all_attrs = [attr for attr in dir(model) if not attr.startswith('_')] + info["all_model_attributes"] = all_attrs + + # Check for specific important attributes + important_attrs = ['adm_in_channels', 'y_dim', 'label_emb', 'time_embed', 'model_channels', + 'context_dim', 'in_channels', 'conv_in', 'num_heads', 'num_res_blocks'] + found_attrs = {attr: hasattr(model, attr) for attr in important_attrs} + info["important_attributes_found"] = found_attrs + # Detect y dimension (pooled embeddings) from various attributes y_dim_candidates = [] @@ -405,4 +416,70 @@ def introspect_diffusion_model(model): elif info["latent_channels"] == 4: info["latent_spatial"] = 128 # SDXL/SD1.5 use 4 channels + # Count parameters to help identify model + param_count = sum(p.numel() for p in model.parameters()) + info["param_count"] = param_count + info["param_count_billions"] = param_count / 1e9 + + # Use parameter count as additional signal + # SDXL: ~2.5-2.6B parameters + # SD1.5: ~900M parameters + # SD3: ~2B parameters + if param_count > 2e9 and info["y_dim"] is None: + # Likely SDXL but failed to detect y_dim + # Try to infer y_dim = 2816 for SDXL + info["y_dim"] = 2816 + info["y_dim_source"] = "inferred_from_param_count_SDXL" + info["has_y_param"] = True + info["architecture"] = "SDXL-like (inferred from 2.5B params)" + + # FINAL CHECK: Test forward pass to definitively determine if y is needed + # This is the most reliable method + if info["y_dim"] is None or not info["has_y_param"]: + print(f" Attempting test forward pass to detect y parameter requirement...") + try: + with torch.no_grad(): + # Create dummy inputs + device = next(model.parameters()).device + dtype = next(model.parameters()).dtype + + test_latent = torch.randn(1, info["latent_channels"], 64, 64, device=device, dtype=dtype) + test_timestep = torch.tensor([999.0], device=device, dtype=dtype) + test_context = torch.randn(1, 77, info["context_dim"], device=device, dtype=dtype) + + # Try without y first + try: + _ = model(test_latent, test_timestep, test_context) + # Success without y - doesn't need it + info["has_y_param"] = False + info["detected_attributes"].append("test_forward_pass_no_y_SUCCESS") + except TypeError as e: + error_str = str(e) + if "must specify y" in error_str or "class-conditional" in error_str: + # Needs y parameter! + info["has_y_param"] = True + # Try to infer y_dim + if info["y_dim"] is None: + # Try common SDXL value + info["y_dim"] = 2816 + info["y_dim_source"] = "inferred_from_forward_pass_error" + info["detected_attributes"].append("test_forward_pass_NEEDS_Y") + elif "unexpected keyword argument" in error_str: + # Model uses different signature + info["detected_attributes"].append(f"test_forward_pass_different_signature: {error_str}") + except Exception as e: + info["detected_attributes"].append(f"test_forward_pass_failed: {str(e)[:100]}") + except Exception as e: + info["detected_attributes"].append(f"test_forward_pass_setup_failed: {str(e)[:100]}") + + # Final architecture determination + if info["y_dim"] == 2816 or info["param_count_billions"] > 2.0: + info["architecture"] = "SDXL-like" + elif info["y_dim"] == 1280: + info["architecture"] = "SDXL-compact" + elif not info["has_y_param"] and info["context_dim"] == 768: + info["architecture"] = "SD1.5-like" + elif info["context_dim"] == 4096: + info["architecture"] = "SD3-like" + return info diff --git a/pyproject.toml b/pyproject.toml index c27952c..c0a8f5b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfy-modelopt" -version = "0.2.0" +version = "0.2.1" description = "NVIDIA ModelOpt integration for ComfyUI - Quantize and optimize diffusion models with INT8/FP8/INT4" readme = "README.md" requires-python = ">=3.10,<3.13" @@ -9,18 +9,6 @@ authors = [ {name = "ComfyUI ModelOpt Contributors"} ] keywords = ["comfyui", "nvidia", "modelopt", "tensorrt", "quantization", "stable-diffusion", "optimization"] -classifiers = [ - "Development Status :: 4 - Beta", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: MIT License", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", -] dependencies = [ "torch>=2.0.0", @@ -73,8 +61,8 @@ target-version = ['py310', 'py311', 'py312'] # ComfyUI Manager metadata display_name = "NVIDIA ModelOpt" description = "Quantize and optimize diffusion models with NVIDIA ModelOpt (INT8/FP8/INT4)" -author = "ComfyUI ModelOpt Contributors" -version = "0.2.0" +author = "enragedantelope" +version = "0.2.1" node_list = [ "ModelOptUNetLoader", "ModelOptQuantizeUNet",