Support auto detecting some zsnr anime checkpoints.

comfyanonymous · Nov 11, 2024 · 8b275ce · 8b275ce
1 parent 2a18e98
commit 8b275ce
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
@@ -197,6 +197,8 @@ def model_type(self, state_dict, prefix=""):
                 self.sampling_settings["sigma_min"] = float(state_dict["edm_vpred.sigma_min"].item())
             return model_base.ModelType.V_PREDICTION_EDM
         elif "v_pred" in state_dict:
+            if "ztsnr" in state_dict: #Some zsnr anime checkpoints
+                self.sampling_settings["zsnr"] = True
             return model_base.ModelType.V_PREDICTION
         else:
             return model_base.ModelType.EPS