Update modules/finetunedvqgan.py

modules/finetunedvqgan.py

@@ -7,8 +7,9 @@ from taming.models.vqgan import GumbelVQ
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class Generator:
-    def __init__(self, config_path, device=device):
+    def __init__(self, config_path, checkpoint_path, device=device):
         self.config_path = config_path
+        self.checkpoint_path = checkpoint_path
         self.device = device
 
     def load_models(self):
@@ -25,7 +26,9 @@ class Generator:
             kl_weight=vq_params.kl_weight,
             temperature_scheduler_config=vq_params.temperature_scheduler_config,
         ).to(self.device)
+        # Load model checkpoints
+        checkpoint = torch.load(self.checkpoint_path, map_location=self.device)
+        # Load the state dictionary into the models
+        model_vaq.load_state_dict(checkpoint, strict=True)
 
-        return model_vaq
-
-
+        return model_vaq