debug

app.py

@@ -140,6 +140,10 @@ class Handler:
         self.vqgan.to(self.device)
         self.vqgan.eval()
 
+        self.optvq = VQModelHF.from_pretrained("BorelTHU/optvq-16x16x4")
+        self.optvq.to(self.device)
+        self.optvq.eval()
+
         self.vae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers")
         self.vae.to(self.device)
         self.vae.eval()
@@ -176,14 +180,19 @@ class Handler:
             revq_rec = revq_rec.contiguous()
             revq_rec = self.preprocesser.inverse(revq_rec)
             revq_rec = self.vae.decode(revq_rec).sample
+            # optvq
+            quant, *_ = self.optvq.encode(img)
+            optvq_rec = self.optvq.decode(quant)
         
         # tensor to PIL image
         img = self.tensor_to_image(img)
         basevq_rec = self.tensor_to_image(basevq_rec)
         vqgan_rec = self.tensor_to_image(vqgan_rec)
         revq_rec = self.tensor_to_image(revq_rec)
+        optvq_rec = self.tensor_to_image(optvq_rec)
         print("Shapes:", img.shape, basevq_rec.shape, vqgan_rec.shape, revq_rec.shape)
-        return img, basevq_rec, vqgan_rec, revq_rec
+        # return img, basevq_rec, vqgan_rec, revq_rec
+        return img, basevq_rec, vqgan_rec, optvq_rec
 
 if __name__ == "__main__":
     # create the model handler