Update app.py

app.py

@@ -79,10 +79,7 @@ def process(
     seed: int,
     tile_diffusion: bool,
     tile_diffusion_size: int,
-    tile_diffusion_stride: int,
-    tile_vae: bool,
-    vae_encoder_tile_size: int,
-    vae_decoder_tile_size: int
+    tile_diffusion_stride: int
 ):
     print(
         f"control image shape={control_img.size}\n"
@@ -91,7 +88,6 @@ def process(
         f"cdf scale={cfg_scale}, steps={steps}, use_color_fix={use_color_fix}\n"
         f"seed={seed}\n"
         f"tile_diffusion={tile_diffusion}, tile_diffusion_size={tile_diffusion_size}, tile_diffusion_stride={tile_diffusion_stride}"
-        f"tile_vae={tile_vae}, vae_encoder_tile_size={vae_encoder_tile_size}, vae_decoder_tile_size={vae_decoder_tile_size}"
     )
     if seed == -1:
         seed = random.randint(0, 2**32 - 1)
@@ -125,7 +121,7 @@ def process(
         shape = (1, 4, height // 8, width // 8)
         x_T = torch.randn(shape, device=model.device, dtype=torch.float32)
 
-        if not tile_diffusion and not tile_vae:
+        if not tile_diffusion:
             samples = sampler.sample_ccsr(
                 steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
                 positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
@@ -133,23 +129,13 @@ def process(
                 color_fix_type="adain" if use_color_fix else "none"
             )
         else:
-            if tile_vae:
-                model._init_tiled_vae(encoder_tile_size=vae_encoder_tile_size, decoder_tile_size=vae_decoder_tile_size)
-            if tile_diffusion:
-                samples = sampler.sample_with_tile_ccsr(
-                    tile_size=tile_diffusion_size, tile_stride=tile_diffusion_stride,
-                    steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
-                    positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
-                    cfg_scale=cfg_scale,
-                    color_fix_type="adain" if use_color_fix else "none"
-                )
-            else:
-                samples = sampler.sample_ccsr(
-                    steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
-                    positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
-                    cfg_scale=cfg_scale,
-                    color_fix_type="adain" if use_color_fix else "none"
-                )
+            samples = sampler.sample_with_tile_ccsr(
+                tile_size=tile_diffusion_size, tile_stride=tile_diffusion_stride,
+                steps=steps, t_max=0.6667, t_min=0.3333, shape=shape, cond_img=control,
+                positive_prompt=positive_prompt, negative_prompt=negative_prompt, x_T=x_T,
+                cfg_scale=cfg_scale,
+                color_fix_type="adain" if use_color_fix else "none"
+            )
 
         x_samples = samples.clamp(0, 1)
         x_samples = (einops.rearrange(x_samples, "b c h w -> b h w c") * 255).cpu().numpy().clip(0, 255).astype(np.uint8)
@@ -195,12 +181,9 @@ with block:
             tile_diffusion = gr.Checkbox(label="Tile diffusion", value=False, info="Enable tiled diffusion for large images.")
             tile_diffusion_size = gr.Slider(label="Tile diffusion size", minimum=512, maximum=1024, value=512, step=256, info="Size of each tile for tiled diffusion.")
             tile_diffusion_stride = gr.Slider(label="Tile diffusion stride", minimum=256, maximum=512, value=256, step=128, info="Stride between tiles for tiled diffusion.")
-            tile_vae = gr.Checkbox(label="Tile VAE", value=True, info="Enable tiled VAE for large images.")
-            vae_encoder_tile_size = gr.Slider(label="Encoder tile size", minimum=512, maximum=5000, value=1024, step=256, info="Size of each tile for the VAE encoder.")
-            vae_decoder_tile_size = gr.Slider(label="Decoder tile size", minimum=64, maximum=512, value=224, step=128, info="Size of each tile for the VAE decoder.")
     
     with gr.Column():
-        result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery")
+        result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery").style(grid=2, height="auto")
 
     inputs = [
         input_image,
@@ -216,9 +199,6 @@ with block:
         tile_diffusion,
         tile_diffusion_size,
         tile_diffusion_stride,
-        tile_vae,
-        vae_encoder_tile_size,
-        vae_decoder_tile_size,
     ]
     run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])