Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+file filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,50 @@
+---
+title: BrushNet
+emoji: ⚡
+colorFrom: yellow
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.50.2
+python_version: 3.9
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+# BrushNet
+
+This repository contains the gradio demo of the paper "BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion"
+
+Keywords: Image Inpainting, Diffusion Models, Image Generation
+
+> [Xuan Ju](https://github.com/juxuan27)<sup>12</sup>, [Xian Liu](https://alvinliu0.github.io/)<sup>12</sup>, [Xintao Wang](https://xinntao.github.io/)<sup>1*</sup>, [Yuxuan Bian](https://scholar.google.com.hk/citations?user=HzemVzoAAAAJ&hl=zh-CN&oi=ao)<sup>2</sup>, [Ying Shan](https://www.linkedin.com/in/YingShanProfile/)<sup>1</sup>, [Qiang Xu](https://cure-lab.github.io/)<sup>2*</sup><br>
+> <sup>1</sup>ARC Lab, Tencent PCG <sup>2</sup>The Chinese University of Hong Kong <sup>*</sup>Corresponding Author
+
+
+<p align="center">
+  <a href="https://tencentarc.github.io/BrushNet/">Project Page</a> |
+  <a href="https://github.com/TencentARC/BrushNet">Code</a> |
+  <a href="https://arxiv.org/abs/2403.06976">Arxiv</a> |
+  <a href="https://forms.gle/9TgMZ8tm49UYsZ9s5">Data</a> |
+  <a href="https://drive.google.com/file/d/1IkEBWcd2Fui2WHcckap4QFPcCI0gkHBh/view">Video</a> |
+</p>
+
+
+## 🤝🏼 Cite Us
+
+```
+@misc{ju2024brushnet,
+  title={BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion}, 
+  author={Xuan Ju and Xian Liu and Xintao Wang and Yuxuan Bian and Ying Shan and Qiang Xu},
+  year={2024},
+  eprint={2403.06976},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV}
+}
+```
+
+
+## 💖 Acknowledgement
+<span id="acknowledgement"></span>
+
+Our code is modified based on [diffusers](https://github.com/huggingface/diffusers), thanks to all the contributors!

app.py

@@ -0,0 +1,341 @@
+##!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+
+# print("Installing correct gradio version...")
+# os.system("pip uninstall -y gradio")
+# os.system("pip install gradio==3.50.0")
+# print("Installing Finished!")
+
+##!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import gradio as gr
+import os
+import cv2
+from PIL import Image
+import numpy as np
+from segment_anything import SamPredictor, sam_model_registry
+import torch
+from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler
+import random
+
+mobile_sam = sam_model_registry['vit_h'](checkpoint='data/ckpt/sam_vit_h_4b8939.pth').to("cuda")
+mobile_sam.eval()
+mobile_predictor = SamPredictor(mobile_sam)
+colors = [(255, 0, 0), (0, 255, 0)]
+markers = [1, 5]
+
+# - - - - - examples  - - - - -  #
+image_examples = [
+    ["examples/brushnet/src/test_image.jpg", "A beautiful cake on the table", "examples/brushnet/src/test_mask.jpg", 0, [], [Image.open("examples/brushnet/src/test_result.png")]],
+    ["examples/brushnet/src/example_1.jpg", "A man in Chinese traditional clothes", "examples/brushnet/src/example_1_mask.jpg", 1, [], [Image.open("examples/brushnet/src/example_1_result.png")]],
+    ["examples/brushnet/src/example_3.jpg", "a cut toy on the table", "examples/brushnet/src/example_3_mask.jpg", 2, [], [Image.open("examples/brushnet/src/example_3_result.png")]],
+    ["examples/brushnet/src/example_4.jpeg", "a car driving in the wild", "examples/brushnet/src/example_4_mask.jpg", 3, [], [Image.open("examples/brushnet/src/example_4_result.png")]],
+    ["examples/brushnet/src/example_5.jpg", "a charming woman wearing dress standing in the dark forest", "examples/brushnet/src/example_5_mask.jpg", 4, [], [Image.open("examples/brushnet/src/example_5_result.png")]],
+]
+
+
+# choose the base model here
+# base_model_path = "data/ckpt/realisticVisionV60B1_v51VAE"
+base_model_path = "runwayml/stable-diffusion-v1-5"
+
+# input brushnet ckpt path
+brushnet_path = "data/ckpt/segmentation_mask_brushnet_ckpt"
+
+brushnet = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch.float16,safety_checker=None)
+pipe = StableDiffusionBrushNetPipeline.from_pretrained(
+    base_model_path, brushnet=brushnet, torch_dtype=torch.float16, low_cpu_mem_usage=False,safety_checker=None
+)
+
+# speed up diffusion process with faster scheduler and memory optimization
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+# remove following line if xformers is not installed or when using Torch 2.0.
+# pipe.enable_xformers_memory_efficient_attention()
+# memory optimization.
+pipe.enable_model_cpu_offload()
+
+def resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img
+
+
+def process(input_image, 
+    original_image, 
+    original_mask, 
+    input_mask, 
+    selected_points, 
+    prompt, 
+    negative_prompt, 
+    blended, 
+    invert_mask, 
+    control_strength, 
+    seed, 
+    randomize_seed, 
+    guidance_scale, 
+    num_inference_steps):
+    if original_image is None:
+        raise gr.Error('Please upload the input image')
+    if (original_mask is None or len(selected_points)==0) and input_mask is None:
+        raise gr.Error("Please click the region where you hope unchanged/changed, or upload a white-black Mask image")
+    
+    # load example image
+    if isinstance(original_image, int):
+        image_name = image_examples[original_image][0]
+        original_image = cv2.imread(image_name)
+        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
+
+    if input_mask is not None:
+        H,W=original_image.shape[:2]
+        original_mask = cv2.resize(input_mask, (W, H))
+    else:
+        original_mask = np.clip(255 - original_mask, 0, 255).astype(np.uint8)
+
+    if invert_mask:
+        original_mask=255-original_mask
+
+    mask = 1.*(original_mask.sum(-1)>255)[:,:,np.newaxis]
+    masked_image = original_image * (1-mask)
+
+    init_image = Image.fromarray(masked_image.astype(np.uint8)).convert("RGB")
+    mask_image = Image.fromarray(original_mask.astype(np.uint8)).convert("RGB")
+    mask_image.save("./mask.png")
+
+    generator = torch.Generator("cuda").manual_seed(random.randint(0,2147483647) if randomize_seed else seed)
+    image_num = 3
+    image = pipe(
+        [prompt]*image_num, 
+        init_image, 
+        mask_image, 
+        num_inference_steps=num_inference_steps, 
+        guidance_scale=guidance_scale,
+        generator=generator,
+        brushnet_conditioning_scale=float(control_strength),
+        negative_prompt=[negative_prompt]*image_num,
+    ).images
+
+    if blended:
+        if control_strength<1.0:
+            raise gr.Error('Using blurred blending with control strength less than 1.0 is not allowed')
+        blended_image=[]
+        # blur, you can adjust the parameters for better performance
+        mask_blurred = cv2.GaussianBlur(mask*255, (21, 21), 0)/255
+        mask_blurred = mask_blurred[:,:,np.newaxis]
+        mask = 1-(1-mask) * (1-mask_blurred)
+        for image_i in image:
+            image_np=np.array(image_i)
+            image_pasted=original_image * (1-mask) + image_np*mask
+
+            image_pasted=image_pasted.astype(image_np.dtype)
+            blended_image.append(Image.fromarray(image_pasted))
+        
+        image=blended_image
+
+    return image
+
+block = gr.Blocks(
+        theme=gr.themes.Soft(
+             radius_size=gr.themes.sizes.radius_none,
+             text_size=gr.themes.sizes.text_md
+         )
+        ).queue()
+with block:
+    with gr.Row():
+        with gr.Column():
+            
+            gr.HTML(f"""
+                    <div style="text-align: center;">
+                        <h1>BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion</h1>
+                        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                            <a href=""></a>
+                            <a href='https://tencentarc.github.io/BrushNet/'><img src='https://img.shields.io/badge/Project_Page-BrushNet-green' alt='Project Page'></a>
+                            <a href='https://arxiv.org/abs/2403.06976'><img src='https://img.shields.io/badge/Paper-Arxiv-blue'></a>
+                        </div>
+                        </br>
+                    </div>
+            """)
+
+
+    with gr.Accordion(label="🧭 Instructions:", open=True, elem_id="accordion"):
+        with gr.Row(equal_height=True):
+            gr.Markdown("""
+            - ⭐️ <b>step1: </b>Upload or select one image from Example
+            - ⭐️ <b>step2: </b>Click on Input-image to select the object to be retained (or upload a white-black Mask image, in which white color indicates the region you want to keep unchanged). You can tick the 'Invert Mask' box to switch region unchanged and change.
+            - ⭐️ <b>step3: </b>Input prompt for generating new contents
+            - ⭐️ <b>step4: </b>Click Run button
+            """)                          
+    with gr.Row():
+        with gr.Column():
+            with gr.Column(elem_id="Input"):
+                with gr.Row():
+                    with gr.Tabs(elem_classes=["feedback"]):
+                        with gr.TabItem("Input Image"):
+                            input_image = gr.Image(type="numpy", label="input",scale=2, height=640)
+                original_image = gr.State(value=None,label="index")
+                original_mask = gr.State(value=None)
+                selected_points = gr.State([],label="select points")
+                with gr.Row(elem_id="Seg"):
+                    radio = gr.Radio(['foreground', 'background'], label='Click to seg: ', value='foreground',scale=2)
+                    undo_button = gr.Button('Undo seg', elem_id="btnSEG",scale=1)
+            prompt = gr.Textbox(label="Prompt", placeholder="Please input your prompt",value='',lines=1)
+            negative_prompt = gr.Text(
+                        label="Negative Prompt",
+                        max_lines=5,
+                        placeholder="Please input your negative prompt",
+                        value='ugly, low quality',lines=1
+                    )
+            with gr.Group():
+                with gr.Row():
+                    blending = gr.Checkbox(label="Blurred Blending", value=False)
+                    invert_mask = gr.Checkbox(label="Invert Mask", value=True)
+            run_button = gr.Button("Run",elem_id="btn")
+            
+            with gr.Accordion("More input params (highly-recommended)", open=False, elem_id="accordion1"):
+                control_strength = gr.Slider(
+                    label="Control Strength: ", show_label=True, minimum=0, maximum=1.1, value=1, step=0.01
+                    )
+                with gr.Group():
+                    seed = gr.Slider(
+                        label="Seed: ", minimum=0, maximum=2147483647, step=1, value=551793204
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                
+                with gr.Group():
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=1,
+                            maximum=12,
+                            step=0.1,
+                            value=12,
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=50,
+                        )
+                with gr.Row(elem_id="Image"):
+                    with gr.Tabs(elem_classes=["feedback1"]):
+                        with gr.TabItem("User-specified Mask Image (Optional)"):
+                            input_mask = gr.Image(type="numpy", label="Mask Image", height=640)
+            
+        with gr.Column():
+            with gr.Tabs(elem_classes=["feedback"]):
+                with gr.TabItem("Outputs"):
+                    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
+    with gr.Row():
+        def process_example(input_image, prompt, input_mask, original_image, selected_points,result_gallery): #
+            return input_image, prompt, input_mask, original_image, [], result_gallery
+        example = gr.Examples(
+            label="Input Example",
+            examples=image_examples,
+            inputs=[input_image, prompt, input_mask, original_image, selected_points,result_gallery],
+            outputs=[input_image, prompt, input_mask, original_image, selected_points],
+            fn=process_example,
+            run_on_click=True,
+            examples_per_page=10
+        )
+
+    # once user upload an image, the original image is stored in `original_image`
+    def store_img(img):
+        # image upload is too slow
+        if min(img.shape[0], img.shape[1]) > 512:
+            img = resize_image(img, 512)
+        if max(img.shape[0], img.shape[1])*1.0/min(img.shape[0], img.shape[1])>2.0:
+            raise gr.Error('image aspect ratio cannot be larger than 2.0')
+        return img, img, [], None  # when new image is uploaded, `selected_points` should be empty
+
+    input_image.upload(
+        store_img,
+        [input_image],
+        [input_image, original_image, selected_points]
+    )
+
+    # user click the image to get points, and show the points on the image
+    def segmentation(img, sel_pix):
+        # online show seg mask
+        points = []
+        labels = []
+        for p, l in sel_pix:
+            points.append(p)
+            labels.append(l)
+        mobile_predictor.set_image(img if isinstance(img, np.ndarray) else np.array(img))
+        with torch.no_grad():
+            masks, _, _ = mobile_predictor.predict(point_coords=np.array(points), point_labels=np.array(labels), multimask_output=False)
+
+        output_mask = np.ones((masks.shape[1], masks.shape[2], 3))*255
+        for i in range(3):
+                output_mask[masks[0] == True, i] = 0.0
+
+        mask_all = np.ones((masks.shape[1], masks.shape[2], 3))
+        color_mask = np.random.random((1, 3)).tolist()[0]
+        for i in range(3):
+            mask_all[masks[0] == True, i] = color_mask[i]
+        masked_img = img / 255 * 0.3 + mask_all * 0.7
+        masked_img = masked_img*255
+        ## draw points
+        for point, label in sel_pix:
+            cv2.drawMarker(masked_img, point, colors[label], markerType=markers[label], markerSize=20, thickness=5)
+        return masked_img, output_mask
+    
+    def get_point(img, sel_pix, point_type, evt: gr.SelectData):
+        if point_type == 'foreground':
+            sel_pix.append((evt.index, 1))   # append the foreground_point
+        elif point_type == 'background':
+            sel_pix.append((evt.index, 0))    # append the background_point
+        else:
+            sel_pix.append((evt.index, 1))    # default foreground_point
+
+        if isinstance(img, int):
+            image_name = image_examples[img][0]
+            img = cv2.imread(image_name)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # online show seg mask
+        masked_img, output_mask = segmentation(img, sel_pix)
+        return masked_img.astype(np.uint8), output_mask
+    
+    input_image.select(
+        get_point,
+        [original_image, selected_points, radio],
+        [input_image, original_mask],
+    )
+
+    # undo the selected point
+    def undo_points(orig_img, sel_pix):
+        # draw points
+        output_mask = None
+        if len(sel_pix) != 0:
+            if isinstance(orig_img, int):   # if orig_img is int, the image if select from examples
+                temp = cv2.imread(image_examples[orig_img][0])
+                temp = cv2.cvtColor(temp, cv2.COLOR_BGR2RGB)
+            else:
+                temp = orig_img.copy()
+            sel_pix.pop()
+            # online show seg mask
+            if len(sel_pix) !=0:
+                temp, output_mask = segmentation(temp, sel_pix)
+            return temp.astype(np.uint8), output_mask
+        else:
+            gr.Error("Nothing to Undo")
+    
+    undo_button.click(
+        undo_points,
+        [original_image, selected_points],
+        [input_image, original_mask]
+    )
+
+    ips=[input_image, original_image, original_mask, input_mask, selected_points, prompt, negative_prompt, blending, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
+    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+
+
+block.launch()

data/ckpt/realisticVisionV60B1_v51VAE/feature_extractor/preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPFeatureExtractor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  }
+}

data/ckpt/realisticVisionV60B1_v51VAE/model_index.json

@@ -0,0 +1,37 @@
+{
+  "_class_name": "StableDiffusionPipeline",
+  "_diffusers_version": "0.26.0.dev0",
+  "feature_extractor": [
+    "transformers",
+    "CLIPFeatureExtractor"
+  ],
+  "image_encoder": [
+    null,
+    null
+  ],
+  "requires_safety_checker": true,
+  "safety_checker": [
+    "stable_diffusion",
+    "StableDiffusionSafetyChecker"
+  ],
+  "scheduler": [
+    "diffusers",
+    "PNDMScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

data/ckpt/realisticVisionV60B1_v51VAE/safety_checker/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "CompVis/stable-diffusion-safety-checker",
+  "architectures": [
+    "StableDiffusionSafetyChecker"
+  ],
+  "initializer_factor": 1.0,
+  "logit_scale_init_value": 2.6592,
+  "model_type": "clip",
+  "projection_dim": 768,
+  "text_config": {
+    "dropout": 0.0,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "model_type": "clip_text_model",
+    "num_attention_heads": 12
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.2",
+  "vision_config": {
+    "dropout": 0.0,
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "model_type": "clip_vision_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "patch_size": 14
+  }
+}

data/ckpt/realisticVisionV60B1_v51VAE/safety_checker/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:afa7ebf10b23008ecbb81111e9cc8818443e1302a0defaa0a6cbd8cfe310b278
+size 1216059369

data/ckpt/realisticVisionV60B1_v51VAE/scheduler/scheduler_config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "PNDMScheduler",
+  "_diffusers_version": "0.26.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}

data/ckpt/realisticVisionV60B1_v51VAE/text_encoder/config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.2",
+  "vocab_size": 49408
+}

data/ckpt/realisticVisionV60B1_v51VAE/text_encoder/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccf6d26500a883b3e7227d2f10292bee7fe078f26cb1afef46eee26252805f0b
+size 492304313

data/ckpt/realisticVisionV60B1_v51VAE/tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

data/ckpt/realisticVisionV60B1_v51VAE/tokenizer/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

data/ckpt/realisticVisionV60B1_v51VAE/unet/config.json

@@ -0,0 +1,67 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.26.0.dev0",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

data/ckpt/realisticVisionV60B1_v51VAE/unet/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8dac118ff5749ca65298f1546ec23d3fec876cded41f9559fcd863e30d969f4
+size 3438354725

data/ckpt/realisticVisionV60B1_v51VAE/vae/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.26.0.dev0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 512,
+  "scaling_factor": 0.18215,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

data/ckpt/realisticVisionV60B1_v51VAE/vae/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9428a7d3cab924439bbba9f2c3e5128d466f9dffd9c12ed7098ae288a5bfe8bd
+size 334707473

data/ckpt/sam_vit_h_4b8939.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e
+size 2564550879

data/ckpt/segmentation_mask_brushnet_ckpt/config.json

@@ -0,0 +1,58 @@
+{
+  "_class_name": "BrushNetModel",
+  "_diffusers_version": "0.27.0.dev0",
+  "_name_or_path": "runs/logs/brushnet_segmask/checkpoint-550000",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "brushnet_conditioning_channel_order": "rgb",
+  "class_embed_type": null,
+  "conditioning_channels": 5,
+  "conditioning_embedding_out_channels": [
+    16,
+    32,
+    96,
+    256
+  ],
+  "cross_attention_dim": 768,
+  "down_block_types": [
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "global_pool_conditions": false,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "MidBlock2D",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_time_scale_shift": "default",
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

data/ckpt/segmentation_mask_brushnet_ckpt/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:066793867083310c25b0d57d300b1ecdc02cb708a155060f020be64afb2ae9c3
+size 2475354520

requirements.txt

@@ -0,0 +1,19 @@
+torch
+torchvision
+torchaudio
+transformers>=4.25.1
+gradio==3.50.0
+ftfy
+tensorboard
+datasets
+Pillow==9.5.0 
+opencv-python 
+imgaug 
+accelerate==0.20.3
+image-reward
+hpsv2
+torchmetrics
+open-clip-torch
+clip
+segment_anything
+git+https://github.com/TencentARC/BrushNet.git