| import os |
| import json |
| import argparse |
|
|
| import torch |
| from diffusers import UNet2DConditionModel |
| from safetensors.torch import save_file |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--src", type=str, default="unet_old/config.json") |
| parser.add_argument("--out", type=str, required=True) |
| parser.add_argument("--layers-per-block", type=str, default=None) |
| parser.add_argument("--transformer-layers", type=str, default=None) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
|
|
| with open(args.src, "r", encoding="utf-8") as f: |
| cfg = json.load(f) |
|
|
| for k in ["_class_name", "_diffusers_version", "_name_or_path"]: |
| cfg.pop(k, None) |
|
|
| if args.layers_per_block is not None: |
| val = args.layers_per_block |
| try: |
| cfg["layers_per_block"] = int(val) |
| except ValueError: |
| cfg["layers_per_block"] = [int(x) for x in val.split(",")] |
| if args.transformer_layers is not None: |
| cfg["transformer_layers_per_block"] = [int(x) for x in args.transformer_layers.split(",")] |
|
|
| torch.manual_seed(args.seed) |
| model = UNet2DConditionModel(**cfg) |
| params = sum(p.numel() for p in model.parameters()) / 1e9 |
| print(f"Created {args.out}: {params:.3f}B params") |
|
|
| os.makedirs(args.out, exist_ok=True) |
| with open(os.path.join(args.out, "config.json"), "w", encoding="utf-8") as f: |
| json.dump({ |
| "_class_name": "UNet2DConditionModel", |
| "_diffusers_version": "0.37.1", |
| "_name_or_path": args.out, |
| **cfg, |
| }, f, indent=2) |
|
|
| save_file(model.state_dict(), os.path.join(args.out, "diffusion_pytorch_model.safetensors")) |
| print("Saved") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
