debug

app.py

@@ -144,7 +144,6 @@ class Handler:
         self.vae.to(self.device)
         self.vae.eval()
         self.preprocesser = load_preprocessor(self.device)
-        # 待修改
         self.revq = ReVQ.from_pretrained("AndyRaoTHU/revq-512T")
         self.revq.to(self.device)
         self.revq.eval()

revq/models/revq.py

@@ -36,7 +36,7 @@ class Viewer:
         x = x.view(B, C, H, W)
         return x
 
-class ReVQ(nn.Module, PyTorchModelHubMixin):
+class ReVQ(PyTorchModelHubMixin, nn.Module):
     def __init__(self, 
             decoder: dict = {},
             quantize: dict = {},
@@ -52,6 +52,23 @@ class ReVQ(nn.Module, PyTorchModelHubMixin):
             checkpoint = torch.load(ckpt_path, map_location="cpu")
             self.decoder.load_state_dict(checkpoint["decoder"])
             self.quantizer.load_state_dict(checkpoint["quantizer"])
+
+    @classmethod
+    def _from_pretrained(cls, repo_id: str, **kwargs):
+        config_path = hf_hub_download(repo_id=repo_id, filename="512T_NC=16384.yaml")
+        ckpt_path = hf_hub_download(repo_id=repo_id, filename="ckpt.pth")
+        
+        full_cfg = OmegaConf.load(config_path)
+        model_cfg = full_cfg.get("model", {})
+        decoder_config = model_cfg.get("decoder", {})
+        quantize_config = model_cfg.get("quantizer", {})
+
+        model = cls(decoder=decoder_config, 
+                    quantize=quantize_config, 
+                    ckpt_path=ckpt_path,
+                    **kwargs)
+
+        return model
     
     def setup_quantizer(self, quantizer_config):
         quantizer = Quantizer(**quantizer_config)
@@ -70,21 +87,4 @@ class ReVQ(nn.Module, PyTorchModelHubMixin):
     def forward(self, x):
         quant = self.quantize(x)
         rec = self.decode(quant)
-        return quant, rec
-    
-    @classmethod
-    def _from_pretrained(cls, repo_id: str, **kwargs):
-        config_path = hf_hub_download(repo_id=repo_id, filename="512T_NC=16384.yaml")
-        ckpt_path = hf_hub_download(repo_id=repo_id, filename="ckpt.pth")
-        
-        full_cfg = OmegaConf.load(config_path)
-        model_cfg = full_cfg.get("model", {})
-        decoder_config = model_cfg.get("decoder", {})
-        quantize_config = model_cfg.get("quantizer", {})
-
-        model = cls(decoder=decoder_config, 
-                    quantize=quantize_config, 
-                    ckpt_path=ckpt_path,
-                    **kwargs)
-
-        return model
+        return quant, rec