Update app.py

app.py

@@ -8,7 +8,6 @@ from omegaconf import OmegaConf
 import subprocess
 from tqdm import tqdm
 import requests
-import spaces
 import einops
 import math
 import random
@@ -62,9 +61,6 @@ load_state_dict(model, ckpt, strict=True)
 model.freeze()
 model.to("cuda")
 
-sampler = SpacedSampler(model, var_type="fixed_small")
-
-@spaces.GPU
 @torch.no_grad()
 def process(
     control_img: Image.Image,
@@ -79,7 +75,10 @@ def process(
     seed: int,
     tile_diffusion: bool,
     tile_diffusion_size: int,
-    tile_diffusion_stride: int
+    tile_diffusion_stride: int,
+    tile_vae: bool,
+    vae_encoder_tile_size: int,
+    vae_decoder_tile_size: int
 ):
     print(
         f"control image shape={control_img.size}\n"
@@ -88,6 +87,7 @@ 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}"
     )
     pl.seed_everything(seed)
 
@@ -118,12 +118,13 @@ def process(
     height, width = control.size(-2), control.size(-1)
     model.control_scales = [strength] * 13
 
+    sampler = SpacedSampler(model, var_type="fixed_small")
     preds = []
     for _ in tqdm(range(num_samples)):
         shape = (1, 4, height // 8, width // 8)
         x_T = torch.randn(shape, device=model.device, dtype=torch.float32)
 
-        if not tile_diffusion:
+        if not tile_diffusion and not tile_vae:
             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,
@@ -131,13 +132,23 @@ def process(
                 color_fix_type="adain" if use_color_fix else "none"
             )
         else:
-            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"
-            )
+            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"
+                )
 
         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)
@@ -182,6 +193,9 @@ with 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.Column():
         result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery")
@@ -200,6 +214,9 @@ 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])