Update all files for BitDance-ImageNet-diffusers

BitDance_B_16x/modeling_autoencoder.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import torch
+from safetensors.torch import load_file as load_safetensors
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from .transformer.qae import VQModel
+except ImportError:  # pragma: no cover
+    from transformer.qae import VQModel
+
+
+class BitDanceImageNetAutoencoder(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self, ddconfig: Dict[str, Any], num_codebooks: int = 4):
+        super().__init__()
+        self.runtime_model = VQModel(ddconfig, num_codebooks)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs):
+        del args, kwargs
+        model_dir = Path(pretrained_model_name_or_path)
+        config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
+        model = cls(ddconfig=config["ddconfig"], num_codebooks=int(config.get("num_codebooks", 4)))
+        state = load_safetensors(model_dir / "diffusion_pytorch_model.safetensors")
+        model.runtime_model.load_state_dict(state, strict=True)
+        model.eval()
+        return model
+
+    def encode(self, x: torch.Tensor):
+        return self.runtime_model.encode(x)
+
+    def decode(self, z: torch.Tensor):
+        return self.runtime_model.decode(z)
+
+    def forward(self, z: torch.Tensor):
+        return self.decode(z)