fix the SD1 AYS model assignment (#77)

MaraScott · Jun 14, 2024 · c97321e · c97321e
1 parent f1db4f1
commit c97321e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/py/nodes/UpscalerRefiner/McBoaty_v2.py b/py/nodes/UpscalerRefiner/McBoaty_v2.py
@@ -107,13 +107,13 @@ def _get_info(self, image_width, image_height, image_divisible_by_8, output_imag
 """]        
 
     @classmethod    
-    def _get_sigmas(self, sigmas_type, model, steps, denoise, scheduler, model_type):
+    def _get_sigmas(self, sigmas_type, model, steps, denoise, scheduler, ays_model_type):
         if sigmas_type == "SDTurboScheduler":
             SigmaScheduler = getattr(comfy_extras.nodes_custom_sampler, sigmas_type)
             sigmas = SigmaScheduler.get_sigmas(SigmaScheduler, model, steps, denoise)[0]
         elif sigmas_type == "AlignYourStepsScheduler":
             SigmaScheduler = AlignYourStepsScheduler
-            sigmas = SigmaScheduler.get_sigmas(SigmaScheduler, model_type, steps, denoise)[0]
+            sigmas = SigmaScheduler.get_sigmas(SigmaScheduler, ays_model_type, steps, denoise)[0]
         else: # BasicScheduler
             SigmaScheduler = getattr(comfy_extras.nodes_custom_sampler, sigmas_type)
             sigmas = SigmaScheduler.get_sigmas(SigmaScheduler, model, scheduler, steps, denoise)[0]
@@ -226,8 +226,8 @@ def fn(self, **kwargs):
         add_noise = True        
         denoise = kwargs.get('denoise', None)        
         sigmas_type = kwargs.get('sigmas_type', None)
-        model_type = kwargs.get('ays_model_type', None)
-        sigmas = self._get_sigmas(sigmas_type, model, steps, denoise, scheduler, model_type)
+        ays_model_type = kwargs.get('ays_model_type', None)
+        sigmas = self._get_sigmas(sigmas_type, model, steps, denoise, scheduler, ays_model_type)
         max_iterations = kwargs.get('running_count', 1)
 
         output_info = [f"No info"]

diff --git a/py/nodes/UpscalerRefiner/McBoaty_v3.py b/py/nodes/UpscalerRefiner/McBoaty_v3.py
@@ -28,15 +28,15 @@
 class UpscalerRefiner_McBoaty_v3():
 
     SIGMAS_TYPES = [
-        "BasicScheduler"
-        , "SDTurboScheduler"
-        , "AlignYourStepsScheduler"
+        'BasicScheduler'
+        , 'SDTurboScheduler'
+        , 'AlignYourStepsScheduler'
     ]    
-    MODEL_TYPE_SIZES = {
-        "SD1.5": 512,
-        "SDXL": 1024,
-        "SD3": 1024,
-        "SVD": 1024,
+    AYS_MODEL_TYPE_SIZES = {
+        'SD1': 512,
+        'SDXL': 1024,
+        'SD3': 1024,
+        'SVD': 1024,
     }
     COLOR_MATCH_METHODS = [   
         'none',
@@ -48,7 +48,7 @@ class UpscalerRefiner_McBoaty_v3():
         'hm-mkl-hm',
     ]
 
-    AYS_MODEL_TYPES = list(MODEL_TYPE_SIZES.keys())
+    AYS_MODEL_TYPES = list(AYS_MODEL_TYPE_SIZES.keys())
 
     INPUTS = {}
     OUTPUTS = {}
@@ -177,15 +177,16 @@ def init(self, **kwargs):
             negative = kwargs.get('negative', None),
             add_noise = True,
             sigmas_type = kwargs.get('sigmas_type', None),
-            model_type = kwargs.get('ays_model_type', None),
+            ays_model_type = kwargs.get('ays_model_type', None),
             steps = kwargs.get('steps', None),
             cfg = kwargs.get('cfg', None),
             denoise = kwargs.get('denoise', None),
         )
+
         self.KSAMPLER.sampler = comfy_extras.nodes_custom_sampler.KSamplerSelect().get_sampler(self.KSAMPLER.sampler_name)[0]
-        self.KSAMPLER.tile_size_sampler = self.MODEL_TYPE_SIZES[self.KSAMPLER.model_type]
-        self.KSAMPLER.sigmas = self._get_sigmas(self.KSAMPLER.sigmas_type, self.KSAMPLER.model, self.KSAMPLER.steps, self.KSAMPLER.denoise, self.KSAMPLER.scheduler, self.KSAMPLER.model_type)
-        self.KSAMPLER.outpaint_sigmas = self._get_sigmas(self.KSAMPLER.sigmas_type, self.KSAMPLER.model, self.KSAMPLER.steps, 1, self.KSAMPLER.scheduler, self.KSAMPLER.model_type)
+        self.KSAMPLER.tile_size_sampler = self.AYS_MODEL_TYPE_SIZES[self.KSAMPLER.ays_model_type]
+        self.KSAMPLER.sigmas = self._get_sigmas(self.KSAMPLER.sigmas_type, self.KSAMPLER.model, self.KSAMPLER.steps, self.KSAMPLER.denoise, self.KSAMPLER.scheduler, self.KSAMPLER.ays_model_type)
+        self.KSAMPLER.outpaint_sigmas = self._get_sigmas(self.KSAMPLER.sigmas_type, self.KSAMPLER.model, self.KSAMPLER.steps, 1, self.KSAMPLER.scheduler, self.KSAMPLER.ays_model_type)
 
         # TODO : make the feather_mask proportional to tile size ?
         # self.PARAMS.feather_mask = self.KSAMPLER.tile_size // 16
@@ -216,13 +217,13 @@ def _get_info(self, image_width, image_height, image_divisible_by_8, output_imag
 """]        
 
     @classmethod    
-    def _get_sigmas(self, sigmas_type, model, steps, denoise, scheduler, model_type):
+    def _get_sigmas(self, sigmas_type, model, steps, denoise, scheduler, ays_model_type):
         if sigmas_type == "SDTurboScheduler":
             SigmaScheduler = getattr(comfy_extras.nodes_custom_sampler, sigmas_type)
             sigmas = SigmaScheduler().get_sigmas(model, steps, denoise)[0]
         elif sigmas_type == "AlignYourStepsScheduler":
             SigmaScheduler = AlignYourStepsScheduler
-            sigmas = SigmaScheduler().get_sigmas(model_type, steps, denoise)[0]
+            sigmas = SigmaScheduler().get_sigmas(ays_model_type, steps, denoise)[0]
         else: # BasicScheduler
             SigmaScheduler = getattr(comfy_extras.nodes_custom_sampler, sigmas_type)
             sigmas = SigmaScheduler().get_sigmas(model, scheduler, steps, denoise)[0]