custom model support in backend

Author: divamgupta

backends/model_converter/.gitignore

@@ -0,0 +1,191 @@
+build_test/
+optimized_stable_diffusion/
+
+
+HF_weights/
+outputs/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+
+
+pretrained_weights/

backends/model_converter/constants.py

@@ -0,0 +1,54 @@
+from fake_torch import fake_torch_load_zipped
+import json
+import numpy as np
+from constants import SD_SHAPES
+import sys 
+
+# python convert_model.py "/Users/divamgupta/Downloads/hollie-mengert.ckpt"  "/Users/divamgupta/Downloads/hollie-mengert.tdict"
+
+# pyinstaller  convert_model.py  --onefile  --noconfirm --clean # build using intel machine so that its cross platform lol
+
+checkpoint_filename = sys.argv[1]
+out_filename = sys.argv[2]
+
+#TODO add MD5s
+
+_HEADER_BYTES  = [42, 10 , 8, 42] + [0]*20
+
+
+s  = 24
+
+torch_weights = fake_torch_load_zipped(open(checkpoint_filename, "rb"))
+keys_info = {}
+out_file = open( out_filename , "wb")
+
+out_file.write(bytes(_HEADER_BYTES))
+
+for k in torch_weights['state_dict']:
+    assert k in SD_SHAPES , k 
+    np_arr = torch_weights['state_dict'][k]
+    key_bytes = np_arr.tobytes()
+    shape = list(np_arr.shape)
+    assert tuple(shape) == SD_SHAPES[k], (k , shape , SD_SHAPES[k] )
+    dtype = str(np_arr.dtype)
+    e = s + len(key_bytes)
+    out_file.write(key_bytes)
+    keys_info[k] = {"start": s , "end" : e , "shape": shape , "dtype" : dtype }
+    s = e
+
+for k in SD_SHAPES:
+    if 'model_ema' in k or 'betas' in k or 'alphas' in k or 'posterior_' in k:
+        continue
+    assert k in keys_info , k
+
+json_start = s
+info_json = bytes( json.dumps(keys_info)  , 'ascii') 
+json_end = s + len(info_json)
+
+out_file.write(info_json)
+
+out_file.seek(5)
+out_file.write(np.array(json_start).astype('long').tobytes())
+
+out_file.seek(14)
+out_file.write(np.array(json_end).astype('long').tobytes())

backends/model_converter/convert_model.py

@@ -0,0 +1,103 @@
+import pickle
+import numpy as np
+import math
+
+def prod(x):
+    return math.prod(x)
+
+def my_unpickle(fb0):
+  key_prelookup = {}
+  class HackTensor:
+    def __new__(cls, *args):
+      #print(args)
+      ident, storage_type, obj_key, location, obj_size = args[0][0:5]
+      assert ident == 'storage'
+
+      assert prod(args[2]) == obj_size
+      ret = np.zeros(args[2], dtype=storage_type)
+      key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
+      return ret
+
+  class HackParameter:
+    def __new__(cls, *args):
+      #print(args)
+      pass
+
+  class Dummy:
+    pass
+
+  class MyPickle(pickle.Unpickler):
+    def find_class(self, module, name):
+      #print(module, name)
+      if name == 'FloatStorage':
+        return np.float32
+      if name == 'LongStorage':
+        return np.int64
+      if name == 'HalfStorage':
+        return np.float16
+      if module == "torch._utils":
+        if name == "_rebuild_tensor_v2":
+          return HackTensor
+        elif name == "_rebuild_parameter":
+          return HackParameter
+      else:
+        try:
+          return pickle.Unpickler.find_class(self, module, name)
+        except Exception:
+          return Dummy
+
+    def persistent_load(self, pid):
+      return pid
+
+  return MyPickle(fb0).load(), key_prelookup
+
+def fake_torch_load_zipped(fb0, load_weights=True):
+  import zipfile
+  with zipfile.ZipFile(fb0, 'r') as myzip:
+    with myzip.open('archive/data.pkl') as myfile:
+      ret = my_unpickle(myfile)
+    if load_weights:
+      for k,v in ret[1].items():
+        with myzip.open(f'archive/data/{k}') as myfile:
+          if v[2].dtype == "object":
+            print(f"issue assigning object on {k}")
+            continue
+          np.copyto(v[2], np.frombuffer(myfile.read(), v[2].dtype).reshape(v[3]))
+  return ret[0]
+
+def fake_torch_load(b0):
+  import io
+  import struct
+
+  # convert it to a file
+  fb0 = io.BytesIO(b0)
+
+  if b0[0:2] == b"\x50\x4b":
+    return fake_torch_load_zipped(fb0)
+
+  # skip three junk pickles
+  pickle.load(fb0)
+  pickle.load(fb0)
+  pickle.load(fb0)
+
+  ret, key_prelookup = my_unpickle(fb0)
+
+  # create key_lookup
+  key_lookup = pickle.load(fb0)
+  key_real = [None] * len(key_lookup)
+  for k,v in key_prelookup.items():
+    key_real[key_lookup.index(k)] = v
+
+  # read in the actual data
+  for storage_type, obj_size, np_array, np_shape, np_strides in key_real:
+    ll = struct.unpack("Q", fb0.read(8))[0]
+    assert ll == obj_size
+    bytes_size = {np.float32: 4, np.int64: 8}[storage_type]
+    mydat = fb0.read(ll * bytes_size)
+    np.copyto(np_array, np.frombuffer(mydat, storage_type).reshape(np_shape))
+
+    # numpy stores its strides in bytes
+    real_strides = tuple([x*bytes_size for x in np_strides])
+    np_array.strides = real_strides
+
+  return ret