Upload 2 files

quantized/combine.py

@@ -1,11 +1,101 @@
-import safetensors.torch  
-  
-# Load both files  
-quantized_blocks = safetensors.torch.load_file("transformer_blocks.safetensors")  
-unquantized_layers = safetensors.torch.load_file("unquantized_layers.safetensors")  
-  
-# Combine the state dicts  
-combined_state_dict = {**quantized_blocks, **unquantized_layers}  
-  
-# Save the combined model  
-safetensors.torch.save_file(combined_state_dict, "combined.safetensors")
+#!/usr/bin/env python
+
+import safetensors.torch
+import os
+import json
+
+def prompt_file(prompt, default):
+    path = input(f"{prompt} [{default}]: ").strip()
+    return path if path else default
+
+# Prompt for input/output files
+quant_file = prompt_file("Enter path to quantized blocks", "transformer_blocks.safetensors")
+unquant_file = prompt_file("Enter path to unquantized layers", "unquantized_layers.safetensors")
+output_file = prompt_file("Enter path to save combined model", "combined.safetensors")
+metadata_file = prompt_file("Enter path to metadata.json (optional)", "metadata.json")
+
+# Validate file existence
+for f in [quant_file, unquant_file]:
+    if not os.path.isfile(f):
+        raise FileNotFoundError(f"File not found: {f}")
+
+# Load state dicts
+quantized_blocks = safetensors.torch.load_file(quant_file)
+unquantized_layers = safetensors.torch.load_file(unquant_file)
+
+# Warn about key overlaps
+overlap = set(quantized_blocks) & set(unquantized_layers)
+if overlap:
+    print(f"⚠️ Warning: Overlapping keys (unquantized will override): {overlap}")
+
+# Merge state dicts
+combined_state_dict = {**quantized_blocks, **unquantized_layers}
+
+# Attempt to load metadata.json
+metadata = {}
+if os.path.isfile(metadata_file):
+    try:
+        with open(metadata_file, "r", encoding="utf-8") as f:
+            metadata = json.load(f)
+        # Convert nested objects to JSON strings as required by safetensors
+        for k, v in metadata.items():
+            if isinstance(v, dict):
+                metadata[k] = json.dumps(v)
+    except Exception as e:
+        print(f"⚠️ Failed to load metadata from {metadata_file}: {e}")
+        print("⏳ Falling back to hardcoded metadata...")
+        metadata = {}  # Will populate below
+
+# Fallback metadata (if file load failed or was missing)
+if not metadata:
+    metadata = {
+        "model_class": "NunchakuFluxTransformer2dModel",
+        "comfy_config": json.dumps({
+            "model_class": "Flux",
+            "model_config": {
+                "axes_dim": [16, 56, 56],
+                "context_in_dim": 4096,
+                "depth": 19,
+                "depth_single_blocks": 38,
+                "disable_unet_model_creation": True,
+                "guidance_embed": True,
+                "hidden_size": 3072,
+                "image_model": "flux",
+                "in_channels": 16,
+                "mlp_ratio": 4.0,
+                "num_heads": 24,
+                "out_channels": 16,
+                "patch_size": 2,
+                "qkv_bias": True,
+                "theta": 10000,
+                "vec_in_dim": 768
+            }
+        }),
+        "quantization_config": json.dumps({
+            "method": "svdquant",
+            "weight": {"dtype": "int4", "scale_dtype": None, "group_size": 64},
+            "activation": {"dtype": "int4", "scale_dtype": None, "group_size": 64}
+        }),
+        "config": json.dumps({
+            "_class_name": "FluxTransformer2DModel",
+            "_diffusers_version": "0.34.0.dev0",
+            "_name_or_path": "../checkpoints/flux-dev/transformer",
+            "attention_head_dim": 128,
+            "axes_dims_rope": [16, 56, 56],
+            "guidance_embeds": True,
+            "in_channels": 64,
+            "joint_attention_dim": 4096,
+            "num_attention_heads": 24,
+            "num_layers": 19,
+            "num_single_layers": 38,
+            "out_channels": None,
+            "patch_size": 1,
+            "pooled_projection_dim": 768
+        })
+    }
+
+# Save the combined safetensors file
+safetensors.torch.save_file(combined_state_dict, output_file, metadata=metadata)
+
+print(f"\n✅ Combined model saved to: {output_file}")
+print(f"ℹ️ Metadata keys included: {', '.join(metadata.keys())}")

quantized/metadata.json

@@ -0,0 +1,45 @@
+{
+  "model_class": "NunchakuFluxTransformer2dModel",
+  "comfy_config": {
+    "model_class": "Flux",
+    "model_config": {
+      "axes_dim": [16, 56, 56],
+      "context_in_dim": 4096,
+      "depth": 19,
+      "depth_single_blocks": 38,
+      "disable_unet_model_creation": true,
+      "guidance_embed": true,
+      "hidden_size": 3072,
+      "image_model": "flux",
+      "in_channels": 16,
+      "mlp_ratio": 4.0,
+      "num_heads": 24,
+      "out_channels": 16,
+      "patch_size": 2,
+      "qkv_bias": true,
+      "theta": 10000,
+      "vec_in_dim": 768
+    }
+  },
+  "quantization_config": {
+    "method": "svdquant",
+    "weight": {"dtype": "int4", "scale_dtype": null, "group_size": 64},
+    "activation": {"dtype": "int4", "scale_dtype": null, "group_size": 64}
+  },
+  "config": {
+    "_class_name": "FluxTransformer2DModel",
+    "_diffusers_version": "0.34.0.dev0",
+    "_name_or_path": "../checkpoints/flux-dev/transformer",
+    "attention_head_dim": 128,
+    "axes_dims_rope": [16, 56, 56],
+    "guidance_embeds": true,
+    "in_channels": 64,
+    "joint_attention_dim": 4096,
+    "num_attention_heads": 24,
+    "num_layers": 19,
+    "num_single_layers": 38,
+    "out_channels": null,
+    "patch_size": 1,
+    "pooled_projection_dim": 768
+  }
+}