"""Helper script to convert diffusion checkpoints to format used by image generator."""

import os

from absl import app
from absl import flags
import requests
import torch as th


_CKPT_PATH = flags.DEFINE_string(
    "ckpt_path", default=None, help="Path to checkpoint file", required=True)
_OUTPUT_PATH = flags.DEFINE_string(
    "output_path", default="bins", help="Output folder path", required=False)

VOCAB_URL = "https://openaipublic.blob.core.windows.net/clip/bpe_simple_vocab_16e6.txt"


def run(ckpt_path, output_path):
  """Converts the checkpoint and saves the result.

  Args:
    ckpt_path: Source checkpoint path
    output_path: Result folder directory
  """
  os.makedirs(output_path, exist_ok=True)
  ckpt = th.load(ckpt_path, map_location="cpu")

  vocab_dest = os.path.join(output_path, os.path.basename(VOCAB_URL))
  if not os.path.exists(vocab_dest):
    with requests.get(VOCAB_URL, stream=True) as response:
      with open(vocab_dest, "wb") as vocab_file:
        for c in response.iter_content(chunk_size=8192):
          vocab_file.write(c)

  for k, v in ckpt["state_dict"].items():
    if "first_stage_model.encoder" in k:
      continue
    if not hasattr(v, "numpy"):
      continue
    output_bin_file = os.path.join(output_path, f"{k}.bin")
    v.numpy().astype("float16").tofile(output_bin_file)


def main(_) -> None:
  ckpt_path = _CKPT_PATH.value
  output_path = _OUTPUT_PATH.value
  run(ckpt_path, output_path)

if __name__ == "__main__":
  app.run(main)