app.py

import gradio as gr
import torch
from PIL import Image, ImageFilter, ImageOps,ImageEnhance
from scipy.ndimage import rank_filter, maximum_filter
import numpy as np
import pathlib
import glob
import os
from diffusers import StableDiffusionControlNetPipeline, DDIMScheduler, AutoencoderKL, ControlNetModel
from ip_adapter import IPAdapter


DESCRIPTION = """# [FilterPrompt](https://arxiv.org/abs/2404.13263): Guiding Imgae Transfer in Diffusion Models
<img id="teaser" alt="teaser" src="https://raw.githubusercontent.com/Meaoxixi/FilterPrompt/gh-pages/resources/teaser.png" />
"""
##################################################################################################################
# 0. Get Pre-Models' Path Ready
##################################################################################################################
vae_model_path = "https://huggingface.co/stabilityai/sd-vae-ft-mse/tree/main"
base_model_path = "https://huggingface.co/runwayml/stable-diffusion-v1-5/tree/main"
image_encoder_path = "https://huggingface.co/h94/IP-Adapter/tree/main/models/image_encoder"
ip_ckpt = "https://huggingface.co/h94/IP-Adapter/tree/main/models/ip-adapter_sd15.bin"
controlnet_softEdge_model_path = "https://huggingface.co/lllyasviel/control_v11p_sd15_softedge/tree/main"
controlnet_depth_model_path = "https://huggingface.co/lllyasviel/control_v11f1p_sd15_depth/tree/main"
# device = "cuda:0"
##################################################################################################################
# 1. load pipeline
##################################################################################################################
torch.cuda.empty_cache()
## 1.1 noise_scheduler
noise_scheduler = DDIMScheduler(
        num_train_timesteps=1000,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        clip_sample=False,
        set_alpha_to_one=False,
        steps_offset=1,
    )
#  1.2 vae
vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
#  1.3 ControlNet
## 1.3.1 load controlnet_softEdge
controlnet_softEdge = ControlNetModel.from_pretrained(controlnet_softEdge_model_path, torch_dtype=torch.float16)
## 1.3.2 load controlnet_depth
controlnet_depth = ControlNetModel.from_pretrained(controlnet_depth_model_path, torch_dtype=torch.float16)
# 1.4 load SD pipeline
pipe_softEdge = StableDiffusionControlNetPipeline.from_pretrained(
        base_model_path,
        controlnet=controlnet_softEdge,
        torch_dtype=torch.float16,
        scheduler=noise_scheduler,
        vae=vae,
        feature_extractor=None,
        safety_checker=None
    )
pipe_depth = StableDiffusionControlNetPipeline.from_pretrained(
        base_model_path,
        controlnet=controlnet_depth,
        torch_dtype=torch.float16,
        scheduler=noise_scheduler,
        vae=vae,
        feature_extractor=None,
        safety_checker=None
    )
print("1 Model loading completed !")
print("##################################################################")
def image_grid(imgs, rows, cols):
    assert len(imgs) == rows * cols
    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols * w, rows * h))
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i % cols * w, i // cols * h))
    return grid
#########################################################################
## funcitions for task 1 : style transfer
#########################################################################
def gaussian_blur(image, blur_radius):
    image = Image.open(image)
    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    return blurred_image

def task1_StyleTransfer(photo, blur_radius, sketch):
    photoImage = Image.open(photo)
    blurPhoto = gaussian_blur(photo, blur_radius)

    Control_factor = 1.2
    IP_factor = 0.6
    ip_model = IPAdapter(pipe_depth, image_encoder_path, ip_ckpt, device, Control_factor=Control_factor, IP_factor=IP_factor)

    depth_image= Image.open(sketch)
    img_array = np.array(depth_image)
    gray_img_array = np.dot(img_array[..., :3], [0.2989, 0.5870, 0.1140])
    # 反相
    inverted_array = 255 - gray_img_array
    gray_img_array = inverted_array.astype(np.uint8)
    processed_image = Image.fromarray(gray_img_array)
    contrast_factor = 2
    enhancer = ImageEnhance.Contrast(processed_image)
    processed_image = enhancer.enhance(contrast_factor)

    images = ip_model.generate(pil_image=photoImage, image=processed_image, num_samples=1, num_inference_steps=30, seed=52)
    original = image_grid(images, 1, 1)
    images = ip_model.generate(pil_image=blurPhoto, image=processed_image, num_samples=1, num_inference_steps=30, seed=52)
    result= image_grid(images, 1, 1)

    return original,result

def task1_test(photo, blur_radius, sketch):
    original = photo
    print(type(original))
    # <class 'str'>
    result = sketch
    return original, result
#########################################################################
## funcitions for task 2 : color transfer
#########################################################################
# todo

#############################################
# Demo
#############################################
theme = gr.themes.Monochrome(primary_hue="blue").set(
        loader_color="#FF0000",
        slider_color="#FF0000",
)

with gr.Blocks(theme=theme) as demo:
    gr.Markdown(DESCRIPTION)

    # 1. 第一个任务Style Transfer的界面代码（青铜器拓本转照片）
    with gr.Group():
        ## 1.1 任务描述
        gr.Markdown(
            """
            ## Case 1: Style transfer
                - In this task, our main goal is to achieve the style transfer from sketch to photo.
                - In the original generation result, the surface of the object has redundant pattern representation from the style image.
                - Next, you can control the Gaussian kernel size of GaussianBlur to weaken the expression of redundant pattern features in the generated results.
            """)
        ## 1.2 输入输出控件布局
        #### 用Column()控制空间在列上的排列关系
        with gr.Row():
            # 第一列
            with gr.Column():
                with gr.Row():
                    ### 1.2.1.1 输入真实照片
                    photo = gr.Image(label="Input photo", type="filepath")
                    print(photo)
                    print(type(photo))
                with gr.Row():
                    ### 1.2.1.2 高斯核控件
                    gaussianKernel = gr.Slider(minimum=0, maximum=8, step=1, value=2, label="Gaussian Blur Radius")
            # 第二列
            with gr.Column():
                with gr.Row():
                    # 1.2.2.1 输入素描图
                    sketch = gr.Image(label="Input sketch", type="filepath")
                    #print(sketch)
                with gr.Row():
                    # 1.2.2.2 按钮：开始生成图片
                    task1Button = gr.Button("Preprocess")
            # 第三列：显示初始的生成图
            with gr.Column():
                with gr.Row():
                    original_result_task1 = gr.Image(label="Original generation result", interactive=False, type="pil")
            # 第四列：显示使用高斯滤波之后的生成图
            with gr.Column():
                result_image_1 = gr.Image(label="Generate results after using GaussianBlur",type="pil")

        ## 1.3 示例图展示
        with gr.Row():
            paths = sorted(pathlib.Path("images/inputExample").glob("*.jpg"))
            gr.Examples(examples=[[path.as_posix()] for path in paths], inputs = sketch)
        with gr.Row():
            gr.Image(value="images/1_gaussian_filter.png", label=" Task example Image", type="filepath")

    # 1. task 1 - style transfer 的界面代码写完了，现在写控件之间交互的逻辑
    task1Button.click(
        fn=task1_StyleTransfer,
        #fn=task1_test,
        inputs=[photo, gaussianKernel, sketch],
        outputs=[original_result_task1, result_image_1],
    )

##################################################################################################################
# 2. run Demo on gradio
##################################################################################################################

if __name__ == "__main__":
    demo.queue(max_size=5).launch()