20f5592077
Co-authored-by: zty624 <164896039+zty624@users.noreply.github.com> Co-authored-by: yhzx233 <16619108+yhzx233@users.noreply.github.com> Co-authored-by: Monalissaa <50078799+Monalissaa@users.noreply.github.com>
62 lines
2.8 KiB
Python
62 lines
2.8 KiB
Python
import argparse
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
def load_shards(directory, index_name):
|
|
index = json.loads((directory / index_name).read_text())
|
|
state = {}
|
|
for filename in sorted(set(index["weight_map"].values())):
|
|
path = (directory / filename).resolve()
|
|
if not path.is_relative_to(directory.resolve()):
|
|
raise ValueError("Shard path escapes the component directory")
|
|
state.update(load_file(path))
|
|
if state.keys() != index["weight_map"].keys():
|
|
raise ValueError("Shard keys do not match the index")
|
|
return state
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("source", type=Path, help="Exported MOSSImagePipeline directory, not a distributed training checkpoint")
|
|
parser.add_argument("models", type=Path, help="Destination ComfyUI models directory")
|
|
parser.add_argument("--name", default="moss-image", help="Output filename prefix for this checkpoint")
|
|
args = parser.parse_args()
|
|
if not args.name or Path(args.name).name != args.name or args.name in (".", ".."):
|
|
parser.error("--name must be a filename prefix, not a path")
|
|
config = json.loads((args.source / "transformer/config.json").read_text())
|
|
if config.get("_class_name") not in ("MOSImageTransformer", "MOSSImageTransformer"):
|
|
raise ValueError(f"Unsupported MOSSImage class: {config.get('_class_name')!r}")
|
|
for key, expected in {"text_dim": 4096, "freq_dim": 256,
|
|
"latent_channels": 128, "rope_base": 10000.0}.items():
|
|
if config.get(key) != expected:
|
|
raise ValueError(f"Unsupported MOSSImage configuration: {key}={config.get(key)!r}")
|
|
outputs = [args.models / f"diffusion_models/{args.name}-dit.safetensors",
|
|
args.models / f"text_encoders/{args.name}-qwen35.safetensors",
|
|
args.models / f"vae/{args.name}-vae.safetensors"]
|
|
for path in outputs:
|
|
if path.exists():
|
|
raise FileExistsError(path)
|
|
for path in outputs:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
for component, index_name, target in [
|
|
("transformer", "diffusion_pytorch_model.safetensors.index.json", outputs[0]),
|
|
("text_encoder", "model.safetensors.index.json", outputs[1]),
|
|
]:
|
|
state = load_shards(args.source / component, index_name)
|
|
# Keep the exported values and dtype; this is packaging, not quantization.
|
|
temporary = target.with_suffix(".partial")
|
|
save_file(state, temporary, metadata={"format": "pt", "model": "moss-image"})
|
|
temporary.rename(target)
|
|
del state
|
|
temporary = outputs[2].with_suffix(".partial")
|
|
shutil.copyfile(args.source / "vae/diffusion_pytorch_model.safetensors", temporary)
|
|
temporary.rename(outputs[2])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|