app.py

##!/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)
pipe = StableDiffusionBrushNetPipeline.from_pretrained(
    base_model_path, brushnet=brushnet, torch_dtype=torch.float16, low_cpu_mem_usage=False
)

# 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

# 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

def process(original_image, input_mask, 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 input_mask is None:
        raise gr.Error("Please upload a white-black Mask image")
    #resizing input image and mask of the object
    original_image = store_img(original_image)
    input_mask = store_img(input_mask)
    
    H, W = original_image.shape[:2]
    original_mask = cv2.resize(input_mask, (W, H))

    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")
    generator = torch.Generator("cuda").manual_seed(random.randint(0, 2147483647) if randomize_seed else seed)
    image = pipe(
        [prompt]*2,
        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]*2,
    ).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 = []
        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

# Create Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            original_image = gr.Image(type="numpy", label="Original Image")
            input_mask = gr.Image(type="numpy", label="Mask Image")
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt", value='ugly, low quality')
            blended = gr.Checkbox(label="Blurred Blending", value=False)
            invert_mask = gr.Checkbox(label="Invert Mask", value=False)
            control_strength = gr.Slider(label="Control Strength", minimum=0, maximum=1.1, value=1, step=0.01)
            seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=551793204)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=12, step=0.1, value=7.5)
            num_inference_steps = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=50, step=1, value=50)
            #selected_points = gr.State([],label="select points")
            run_button = gr.Button("Run")
        
        with gr.Column():
            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True)
    
    inputs = [original_image, input_mask, prompt, negative_prompt, blended, invert_mask, control_strength, seed, randomize_seed, guidance_scale, num_inference_steps]
    run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])

demo.queue(concurrency_count=1, api_open=True)
demo.launch(show_api=True, enable_queue=True, show_error=True)