fix bug

2024-04-09 14:02:28 +08:00
parent ee7fe888f2
commit 8a9b791e3f
4 changed files with 25 additions and 12 deletions
@@ -395,7 +395,7 @@ class DiffusionInference():
            return self.first_stage_model['paras']['scale_factor'] * z

    def decode_first_stage(self, z):
-        _, dtype = self.get_function_info(self.first_stage_model, 'encode')
+        _, dtype = self.get_function_info(self.first_stage_model, 'decode')
        with torch.autocast('cuda',
                            enabled=dtype == 'float16',
                            dtype=getattr(torch, dtype)):