fix not working apply_late_temperature

app.py

@@ -48,6 +48,7 @@ def infer(
         turn_off_cfg_start_si=turn_off_cfg_start_si,
         turn_on_cfg_start_si=turn_on_cfg_start_si,
         seed=seed,
+        apply_late_temperature=apply_late_temperature,
         last_scale_temp=last_scale_temp,
     )[0]
 
@@ -163,8 +164,6 @@ with gr.Blocks(css=css) as demo:
                     step=0.1,
                     value=1,
                 )
-                if not apply_late_temperature:
-                    last_scale_temp = None
 
 
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True)# cache_mode="lazy")
@@ -183,6 +182,7 @@ with gr.Blocks(css=css) as demo:
             smooth_start_si,
             turn_off_cfg_start_si,
             more_diverse,
+            apply_late_temperature,
             last_scale_temp,
         ],
         outputs=[result, seed],

models/pipeline.py

@@ -93,6 +93,7 @@ class SwittiPipeline:
         turn_off_cfg_start_si: int = 10,
         turn_on_cfg_start_si: int = 0,
         image_size: tuple[int, int] = (512, 512),
+        apply_late_temperature: bool = False,
         last_scale_temp: None | float = None,
     ) -> torch.Tensor | list[PILImage]:
         """
@@ -106,6 +107,8 @@ class SwittiPipeline:
         :param more_smooth: sampling using gumbel softmax; only used in visualization, not used in FID/IS benchmarking
         :return: if return_pil: list of PIL Images, else: torch.tensor (B, 3, H, W) in [0, 1]
         """
+        if not apply_late_temperature:
+            last_scale_temp = None
         assert not self.switti.training
         switti = self.switti
         vae = self.vae