Update app.py

app.py

@@ -79,18 +79,14 @@ 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"
           f"num_samples={num_samples}, sr_scale={sr_scale}, strength={strength}\n"
           f"positive_prompt='{positive_prompt}', negative_prompt='{negative_prompt}'\n"
           f"cfg 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}")
+          f"tile_diffusion={tile_diffusion}, tile_diffusion_size={tile_diffusion_size}, tile_diffusion_stride={tile_diffusion_stride}")
     
     pl.seed_everything(seed)
 
@@ -127,7 +123,7 @@ def process(
         shape = (1, 4, height // 8, width // 8)
         x_T = torch.randn(shape, device=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,
@@ -135,23 +131,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)
@@ -161,9 +147,13 @@ def process(
 
     return preds
 
-def update_output_resolution(image, scale):
+def update_output_resolution(image, scale_choice, custom_scale):
     if image is not None:
         width, height = image.size
+        if scale_choice == "Custom":
+            scale = custom_scale
+        else:
+            scale = float(scale_choice.split()[-1].strip("()x"))
         return f"Current resolution: {width}x{height}. Output resolution: {int(width*scale)}x{int(height*scale)}"
     return "Upload an image to see the output resolution"
 
@@ -233,9 +223,6 @@ with gr.Blocks(css=css) as block:
                 tile_diffusion = gr.Checkbox(label="Tile diffusion", value=False)
                 tile_diffusion_size = gr.Slider(label="Tile diffusion size", minimum=512, maximum=1024, value=512, step=256)
                 tile_diffusion_stride = gr.Slider(label="Tile diffusion stride", minimum=256, maximum=512, value=256, step=128)
-                tile_vae = gr.Checkbox(label="Tile VAE", value=True)
-                vae_encoder_tile_size = gr.Slider(label="Encoder tile size", minimum=512, maximum=5000, value=1024, step=256)
-                vae_decoder_tile_size = gr.Slider(label="Decoder tile size", minimum=64, maximum=512, value=224, step=128)
     
     with gr.Row():
         result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery", elem_classes="output-gallery")
@@ -253,10 +240,10 @@ with gr.Blocks(css=css) as block:
     inputs = [
         input_image, num_samples, sr_scale, strength, positive_prompt, negative_prompt,
         cfg_scale, steps, use_color_fix, seed, tile_diffusion, tile_diffusion_size,
-        tile_diffusion_stride, tile_vae, vae_encoder_tile_size, vae_decoder_tile_size,
+        tile_diffusion_stride
     ]
     run_button.click(
-        fn=lambda *args: process(*args[:1], args[1], get_scale_value(args[2], args[-1]), *args[3:-1]),
+        fn=lambda *args: process(*args[:1], args[1], get_scale_value(args[2], args[-1]), *args[3:]),
         inputs=inputs + [custom_scale],
         outputs=[result_gallery]
     )
@@ -269,12 +256,17 @@ with gr.Blocks(css=css) as block:
 
     input_image.change(
         update_output_resolution,
-        inputs=[input_image, sr_scale],
+        inputs=[input_image, sr_scale, custom_scale],
         outputs=[output_resolution]
     )
     sr_scale.change(
         update_output_resolution,
-        inputs=[input_image, sr_scale],
+        inputs=[input_image, sr_scale, custom_scale],
+        outputs=[output_resolution]
+    )
+    custom_scale.change(
+        update_output_resolution,
+        inputs=[input_image, sr_scale, custom_scale],
         outputs=[output_resolution]
     )