update

app.py

@@ -0,0 +1,197 @@
+from argparse import Namespace
+from glob import glob
+import yaml
+import os
+
+import gradio as gr
+import torch
+import torchvision
+import safetensors
+from diffusers import AutoencoderKL
+from peft import get_peft_model, LoraConfig, set_peft_model_state_dict
+from huggingface_hub import snapshot_download
+
+pretrained_model_path = snapshot_download(repo_id="revp2024/revp-censorship")
+with open(glob(os.path.join(pretrained_model_path, 'hparams.yml'), recursive=True)[0]) as f:
+    args = Namespace(**yaml.safe_load(f))
+
+def prepare_model():
+    print('Loading model ...')
+    vae_lora_config = LoraConfig(
+        r=args.rank,
+        lora_alpha=args.rank,
+        init_lora_weights="gaussian",
+        target_modules=["conv", "conv1", "conv2",
+                        "to_q", "to_k", "to_v", "to_out.0"],
+    )
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae"
+    )
+    vae = get_peft_model(vae, vae_lora_config)
+    lora_weights_path = os.path.join(pretrained_model_path, f"pytorch_lora_weights.safetensors")
+    state_dict = {}
+    with safetensors.torch.safe_open(lora_weights_path, framework="pt", device="cpu") as f:
+       for key in f.keys():
+           state_dict[key] = f.get_tensor(key)
+
+    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')
+
+
+@torch.no_grad()
+def add_censorship(input_image, mode, pixelation_block_size, blur_kernel_size, soft_edges, soft_edge_kernel_size):
+    background, layers, _ = input_image.values()
+    input_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
+
+    H, W = input_images.shape[-2:]
+    if H > 1024 or W > 1024:
+        H_t, W_t = H, W
+        if H > W:
+            H, W = 1024, int(1024 * W_t / H_t)
+        else:
+            H, W = int(1024 * H_t / W_t), 1024
+    H_q8 = (H // 8) * 8
+    W_q8 = (W // 8) * 8
+    input_images = torch.nn.functional.interpolate(input_images, (H_q8, W_q8), mode='bilinear')
+    mask = torch.nn.functional.interpolate(mask, (H_q8, W_q8))
+    if soft_edges:
+        mask = torchvision.transforms.functional.gaussian_blur(mask, soft_edge_kernel_size)[0][0]
+
+    input_images = input_images.to(vae.device)
+
+    if mode == 'Pixelation':
+        censored = torch.nn.functional.avg_pool2d(
+            input_images, pixelation_block_size)
+        censored = torch.nn.functional.interpolate(censored, input_images.shape[-2:])
+    elif mode == 'Gaussian blur':
+        censored = torchvision.transforms.functional.gaussian_blur(
+            input_images, blur_kernel_size)
+    elif mode == 'Black':
+        censored = torch.zeros_like(input_images)
+    else:
+        raise ValueError("censor_mode has to be either `pixelation' or `gaussian_blur'")
+    
+    mask = mask.to(input_images.device)
+    censored_images = input_images * (1 - mask) + censored * mask
+    censored_images *= 255
+
+    input_images = input_images * 2 - 1
+    with vae.disable_adapter():
+        latents = vae.encode(input_images).latent_dist.mean
+    images = vae.decode(latents, return_dict=False)[0]
+
+    # denormalize
+    images = images / 2 + 0.5
+    images *= 255
+
+    residuals = (images - censored_images).clamp(-args.budget, args.budget)
+    images = (censored_images + residuals).clamp(0, 255).to(torch.uint8)
+
+    gr.Info("Try to donwload/copy the censored image to the `Remove censorsip' tab")
+    return images[0].permute(1, 2, 0).cpu().numpy()
+
+@torch.no_grad()
+def remove_censorship(input_image, 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]
+    # denormalize
+    images = images / 2 + 0.5
+    images *= 255
+    images = images.clamp(0, 255).to(torch.uint8)
+    return images[0].permute(1, 2, 0).cpu().numpy()
+
+# @@@@@@@ Start of the program @@@@@@@@
+
+vae = prepare_model()
+
+css = '''
+.my-disabled {
+    background-color: #eee;
+}
+.my-disabled input {
+    background-color: #eee;
+}
+'''
+with gr.Blocks(css=css) as demo:
+    gr.Markdown('# ReVP: Reversible Visual Processing with Latent Models')
+    with gr.Tab('Add censorship'):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.ImageEditor(brush=gr.Brush(default_size=100))
+                with gr.Accordion('Options', open=False) as options_accord:
+                    mode = gr.Radio(label='Mode', choices=['Pixelation', 'Gaussian blur', 'Black'],
+                                    value='Pixelation', interactive=True)
+                    pixelation_block_size = gr.Slider(label='Block size', minimum=10, maximum=40, value=25, step=1, interactive=True)
+                    blur_kernel_size = gr.Slider(label='Blur kernel size', minimum=21, maximum=151,  value=85, step=2, interactive=True, visible=False)
+                    def change_mode(mode):
+                        if mode == 'Gaussian blur':
+                            return gr.Slider(visible=False), gr.Slider(visible=True), gr.Accordion(open=True)
+                        elif mode == 'Pixelation':
+                            return gr.Slider(visible=True), gr.Slider(visible=False), gr.Accordion(open=True)
+                        elif mode == 'Black':
+                            return gr.Slider(visible=False), gr.Slider(visible=False), gr.Accordion(open=True)
+                        else:
+                            raise NotImplementedError
+                    mode.select(change_mode, mode, [pixelation_block_size, blur_kernel_size, options_accord])
+                    with gr.Row(variant='panel'):
+                        soft_edges = gr.Checkbox(label='Soft edges', value=True, interactive=True, scale=1)
+                        soft_edge_kernel_size = gr.Slider(label='Soft edge kernel size', minimum=21, maximum=49,  value=35, step=2, interactive=True, visible=True, scale=2)
+                    def change_soft_edges(soft_edges):
+                        return gr.Slider(visible=True if soft_edges else False), gr.Accordion(open=True)
+                    soft_edges.change(change_soft_edges, soft_edges, [soft_edge_kernel_size, options_accord])
+                submit_btn = gr.Button('Submit')
+            output_image = gr.Image(label='Censored', show_download_button=True)
+
+        submit_btn.click(
+            fn=add_censorship,
+            inputs=[input_image, mode, pixelation_block_size, blur_kernel_size, soft_edges, soft_edge_kernel_size],
+            outputs=output_image
+        )
+
+    with gr.Tab('Remove censorship'):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.ImageEditor()
+                with gr.Accordion('Manual cropping', open=False):
+                    with gr.Row():
+                        with gr.Row():
+                            x1 = gr.Number(value=0, label='x1')
+                            y1 = gr.Number(value=0, label='y1')
+                        with gr.Row():
+                            x2_ = gr.Number(value=10000, label='x2', interactive=False, elem_classes='my-disabled')
+                            y1_ = gr.Number(value=0, label='y1', interactive=False, elem_classes='my-disabled')
+                    with gr.Row():
+                        with gr.Row():
+                            x1_ =gr.Number(value=0, label='x1', elem_classes='my-disabled')
+                            y2_ = gr.Number(value=10000, label='y2', elem_classes='my-disabled')
+                        with gr.Row():
+                            x2 = gr.Number(value=10000, label='x2')
+                            y2 = gr.Number(value=10000, label='y2')
+                submit_btn = gr.Button('Submit')
+            output_image = gr.Image(label='Uncensored')
+
+        submit_btn.click(
+            fn=remove_censorship,
+            inputs=[input_image, x1, y1, x2, y2],
+            outputs=output_image
+        )
+
+        # sync coordinate on changed
+        x1.change(lambda x : x, x1, x1_)
+        x2.change(lambda x : x, x2, x2_)
+        y1.change(lambda x : x, y1, y1_)
+        y2.change(lambda x : x, y2, y2_)
+
+if __name__ == '__main__':
+    demo.queue(4)
+    demo.launch()
+

requirements.txt

@@ -0,0 +1,12 @@
+torch
+torchvision
+huggingface_hub
+accelerate
+transformers
+datasets
+diffusers
+peft
+safetensors
+gradio==4.26.0
+pyyaml
+