import os
import glob
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file, save_file

def merge_and_cast_qwen3_fp8():
    repo_id = "prithivMLmods/Qwen3-VL-4B-Instruct-abliterated-v1"
    output_filename = "qwen3vl_4b_abliterated_fp8.safetensors"

    print(f"Pulling safetensors from {repo_id}...")
    # Downloads only the safetensors to bypass heavy unused binaries
    model_path = snapshot_download(repo_id=repo_id, allow_patterns=["*.safetensors"])

    shard_files = sorted(glob.glob(os.path.join(model_path, "*.safetensors")))
    print(f"Found {len(shard_files)} shards. Merging and casting to FP8...")

    merged_state_dict = {}
    for file in shard_files:
        print(f"Processing {os.path.basename(file)}...")
        shard_dict = load_file(file)

        for key, tensor in shard_dict.items():
            # Cast standard floating point weights down to fp8_e4m3fn
            if tensor.is_floating_point():
                merged_state_dict[key] = tensor.to(torch.float8_e4m3fn)
            else:
                # Keep any integer/index tensors exactly as they are
                merged_state_dict[key] = tensor

    print(f"Saving merged FP8 weights to {output_filename}...")
    save_file(merged_state_dict, output_filename)
    print(f"Done. Place '{output_filename}' into ComfyUI/models/text_encoders/")

if __name__ == "__main__":
    merge_and_cast_qwen3_fp8()