add consistency decoder

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 import torch
 import torchvision
 import safetensors
-from diffusers import AutoencoderKL
+from diffusers import AutoencoderKL, ConsistencyDecoderVAE
 from peft import get_peft_model, LoraConfig, set_peft_model_state_dict
 from huggingface_hub import snapshot_download
 
@@ -43,7 +43,10 @@ def prepare_model():
     set_peft_model_state_dict(vae, state_dict)
 
     print('Done.')
-    return vae.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
+    cd_vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
+    vae = vae.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
+    cd_vae = cd_vae.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
+    return vae, cd_vae
 
 
 @spaces.GPU
@@ -102,15 +105,18 @@ def add_censorship(input_image, mode, pixelation_block_size, blur_kernel_size, s
 
 @spaces.GPU
 @torch.no_grad()
-def remove_censorship(input_image, x1, y1, x2, y2):
+def remove_censorship(input_image, use_cd, x1, y1, x2, y2):
     background, layers, _ = input_image.values()
     images = torch.from_numpy(background).permute(2, 0, 1)[None, :3] / 255
     mask = torch.from_numpy(layers[0]).permute(2, 0, 1)[None, -1:] / 255
     images = images * (1 - mask)
     images = images[..., y1:y2, x1:x2]
     latents = vae.encode((images * 2 - 1).to(vae.device)).latent_dist.mean
-    with vae.disable_adapter():
-        images = vae.decode(latents, return_dict=False)[0]
+    if use_cd:
+        images = cd_vae.decode(latents.to(cd_vae.dtype), return_dict=False)[0]
+    else:
+        with vae.disable_adapter():
+            images = vae.decode(latents, return_dict=False)[0]
     # denormalize
     images = images / 2 + 0.5
     images *= 255
@@ -119,7 +125,7 @@ def remove_censorship(input_image, x1, y1, x2, y2):
 
 # @@@@@@@ Start of the program @@@@@@@@
 
-vae = prepare_model()
+vae, cd_vae = prepare_model()
 
 css = '''
 .my-disabled {
@@ -177,6 +183,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 input_image = gr.ImageEditor()
+                use_cd = gr.Checkbox(label='Use Consistency Decoder (slower)')
                 with gr.Accordion('Manual cropping', open=False):
                     with gr.Row():
                         with gr.Row():
@@ -197,13 +204,13 @@ with gr.Blocks(css=css) as demo:
 
         submit_btn.click(
             fn=remove_censorship,
-            inputs=[input_image, x1, y1, x2, y2],
+            inputs=[input_image, use_cd, x1, y1, x2, y2],
             outputs=output_image
         )
         gr.Examples(
             examples=remove_censor_examples,
             fn=remove_censorship,
-            inputs=[input_image, x1, y1, x2, y2],
+            inputs=[input_image, use_cd, x1, y1, x2, y2],
             outputs=output_image,
             cache_examples=False,
         )

examples/remove_censorship.yaml

@@ -1,14 +1,17 @@
 - - examples/images/processed/car.png
+  - false
   - 0
   - 0
   - 10000
   - 10000
 - - examples/images/processed/obama.png
+  - false
   - 0
   - 0
   - 10000
   - 10000
 - - examples/images/processed/steam-clock.png
+  - false
   - 0
   - 0
   - 10000