import torch
import sys
import os
import json

# Add ComfyUI to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from safetensors import safe_open
from safetensors.torch import save_file
import comfy.model_management
import comfy.sd
import comfy.quant_ops

# --- CONFIG ---
INPUT_FILE = "/your model/weightst.safetensors"
OUTPUT_FILE = "/put something here/"
# FIXED: Set this to krea2 so ComfyUI uses the correct model architecture!
TARGET_ARCH = "krea2" 

# ---------------------------------------------------------
# 1. Read & Clean Metadata
# ---------------------------------------------------------
print("Reading metadata header...")
with safe_open(INPUT_FILE, framework="pt") as f:
    metadata = f.metadata() or {}

# Remove any quantization tags so Star Nodes/ComfyUI doesn't get confused
keys_to_remove = ['quantization', 'quant_method', 'unet_dtype', 'dtype']
for key in list(metadata.keys()):
    if key in keys_to_remove or 'quant' in key.lower():
        metadata.pop(key)

# Force architecture tags and BF16 dtype
metadata["architecture"] = TARGET_ARCH
metadata["unet_type"] = TARGET_ARCH
metadata["transformer_type"] = TARGET_ARCH
metadata["dtype"] = "bf16"
print("Cleaned Metadata:", json.dumps(metadata, indent=2))

# ---------------------------------------------------------
# 2. Load & Patch Model
# ---------------------------------------------------------
print(f"\nLoading {INPUT_FILE}...")
try:
    model_patcher = comfy.sd.load_diffusion_model(INPUT_FILE)
except Exception as e:
    print(f"Failed to load: {e}")
    sys.exit(1)

print("Patching model to instantiate ConvRot objects...")
try:
    model_patcher.patch_model()
except Exception as e:
    print(f"Failed to patch model: {e}")
    sys.exit(1)

# ---------------------------------------------------------
# 3. Extract & Dequantize Directly from Model Modules
# ---------------------------------------------------------
new_sd = {}
dequant_count = 0

# We get the raw state dict to know what keys to skip (like .comfy_quant and .weight_scale)
raw_sd = model_patcher.model_state_dict()
quantized_bases = set()
for k in raw_sd.keys():
    if k.endswith(".comfy_quant"):
        quantized_bases.add(k[:-len(".comfy_quant")])

print("Extracting and dequantizing weights from model tree...")

def process_module(module, prefix=""):
    global dequant_count
    for name, child in module.named_children():
        child_prefix = f"{prefix}.{name}" if prefix else name
        
        if hasattr(child, 'weight'):
            weight_key = child_prefix + ".weight"
            
            # If it's a QuantizedTensor, we MUST use the official dequantize method
            # This properly reverses ConvRot rotation and applies scales
            if isinstance(child.weight, comfy.quant_ops.QuantizedTensor):
                try:
                    dq = child.weight.dequantize()
                    # CAST TO BFLOAT16 to prevent NaN/Black image issues
                    new_sd[weight_key] = dq.to(torch.bfloat16).contiguous()
                    dequant_count += 1
                except Exception as e:
                    print(f"  ERROR dequantizing {weight_key}: {e}")
            elif isinstance(child.weight, torch.Tensor):
                new_sd[weight_key] = child.weight.to(torch.bfloat16).contiguous()
                
        if hasattr(child, 'bias') and child.bias is not None:
            bias_key = child_prefix + ".bias"
            new_sd[bias_key] = child.bias.to(torch.bfloat16).contiguous()
            
        # Recurse into children
        process_module(child, child_prefix)

# Start recursion
process_module(model_patcher.model)

# Copy any remaining non-weight, non-scale, non-metadata tensors (like norms)
for k, v in raw_sd.items():
    if k.endswith(".comfy_quant"): continue
    if k.endswith(".weight_scale") and k[:-len(".weight_scale")] in quantized_bases: continue
    if k in new_sd: continue
    
    if isinstance(v, torch.Tensor):
        if v.dtype == torch.int8:
            print(f"  WARNING: Raw int8 found at {k} without dequantize method. Skipping.")
            continue
        new_sd[k] = v.to(torch.bfloat16).contiguous()

# Cleanup VRAM
try:
    model_patcher.unpatch_model()
except:
    pass

# ---------------------------------------------------------
# 4. Add "model." Prefix & Save
# ---------------------------------------------------------
final_sd = {}
for k, v in new_sd.items():
    # ComfyUI strips "model." in memory, we must add it back for safetensors
    new_key = "model." + k if not k.startswith("model.") else k
    final_sd[new_key] = v

print(f"\nDequantized: {dequant_count} tensors")
print(f"Saving clean BF16 model to {OUTPUT_FILE}...")
save_file(final_sd, OUTPUT_FILE, metadata=metadata)
print("SUCCESS! File is a standard, unscrambled BF16 Krea2 model.")