add assertion for the guidance_scale

mit-han-lab · May 5, 2024 · c71cba0 · c71cba0
1 parent f1f45fb
commit c71cba0
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/distrifuser/pipelines.py b/distrifuser/pipelines.py
@@ -49,6 +49,11 @@ def __call__(self, *args, **kwargs):
         assert "height" not in kwargs, "height should not be in kwargs"
         assert "width" not in kwargs, "width should not be in kwargs"
         config = self.distri_config
+        if not config.do_classifier_free_guidance:
+            if "guidance_scale" not in kwargs:
+                kwargs["guidance_scale"] = 1
+            else:
+                assert kwargs["guidance_scale"] == 1
         self.pipeline.unet.set_counter(0)
         return self.pipeline(height=config.height, width=config.width, *args, **kwargs)
 
@@ -202,6 +207,11 @@ def __call__(self, *args, **kwargs):
         assert "height" not in kwargs, "height should not be in kwargs"
         assert "width" not in kwargs, "width should not be in kwargs"
         config = self.distri_config
+        if not config.do_classifier_free_guidance:
+            if not "guidance_scale" not in kwargs:
+                kwargs["guidance_scale"] = 1
+            else:
+                assert kwargs["guidance_scale"] == 1
         self.pipeline.unet.set_counter(0)
         return self.pipeline(height=config.height, width=config.width, *args, **kwargs)