remove apply_late_temperature checkbox

app.py

@@ -28,8 +28,7 @@ def infer(
     smooth_start_si=2,
     turn_off_cfg_start_si=10,
     more_diverse=True,
-    apply_late_temperature=False,
-    last_scale_temp=None,
+    last_scale_temp=1,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -49,7 +48,6 @@ 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]
 
@@ -141,17 +139,15 @@ with gr.Blocks(css=css) as demo:
                 )
             with gr.Row():
                 more_diverse = gr.Checkbox(label="More diverse", value=True)
-                apply_late_temperature = gr.Checkbox(label="Temperature after disabling CFG", value=False)
                 last_scale_temp = gr.Slider(
-                    label="Late temperature value",
+                    label="Temperature after disabling CFG",
                     minimum=0.1,
                     maximum=10,
                     step=0.1,
                     value=0.1,
                 )
 
-
-        gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True)# cache_mode="lazy")
+        gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)# cache_mode="lazy")
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -167,7 +163,6 @@ 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,8 +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,
+        last_scale_temp: float = 1.,
     ) -> torch.Tensor | list[PILImage]:
         """
         only used for inference, on autoregressive mode
@@ -107,8 +106,6 @@ 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
@@ -200,7 +197,7 @@ class SwittiPipeline:
                 # default const cfg
                 t = cfg
                 logits_BlV = (1 + t) * logits_BlV[:B] - t * logits_BlV[B:]
-            elif last_scale_temp is not None:
+            else:
                 logits_BlV = logits_BlV / last_scale_temp
 
             if apply_smooth and si >= smooth_start_si: