Upload gemma_safetensors.py with huggingface_hub

gemma_safetensors.py

@@ -0,0 +1,45 @@
+import os
+import json
+import torch
+from safetensors.torch import load_file, save_file
+from huggingface_hub import snapshot_download
+
+# 1. Download official model shards from Hugging Face
+model_id = "google/gemma-4-31B-it"
+local_dir = "./gemma-4-31B-it"
+print("Downloading model shards...")
+snapshot_download(repo_id=model_id, local_dir=local_dir, allow_patterns=["*.safetensors", "tokenizer.json"])
+
+# 2. Merge shards and remap key prefixes to match gemma4.py
+state_dict = {}
+print("Remapping keys...")
+for filename in sorted(os.listdir(local_dir)):
+    if filename.endswith(".safetensors"):
+        shard_path = os.path.join(local_dir, filename)
+        shard = load_file(shard_path)
+        for key, tensor in shard.items():
+            new_key = key
+            # Strip Hugging Face wrappers
+            if new_key.startswith("model.language_model.model."):
+                new_key = "model." + new_key[len("model.language_model.model."):]
+            elif new_key.startswith("model.language_model."):
+                new_key = "model." + new_key[len("model.language_model."):]
+            elif new_key.startswith("model.vision_tower."):
+                new_key = "vision_model." + new_key[len("model.vision_tower."):]
+            elif new_key.startswith("model.multi_modal_projector."):
+                new_key = "multi_modal_projector." + new_key[len("model.multi_modal_projector."):]
+
+            state_dict[new_key] = tensor
+
+# 3. Embed tokenizer.json in safetensors metadata header
+metadata = {}
+tokenizer_path = os.path.join(local_dir, "tokenizer.json")
+if os.path.exists(tokenizer_path):
+    with open(tokenizer_path, "r", encoding="utf-8") as f:
+        metadata["tokenizer_json"] = f.read()
+
+# 4. Save unified ComfyUI safetensors file
+output_path = "gemma4_31b_it_bf16.safetensors"
+print(f"Saving merged model to {output_path}...")
+save_file(state_dict, output_path, metadata=metadata)
+print("Conversion complete!")