svjack/ControlNet-Pose-Chinese

1from diffusers import utils2from diffusers.utils import deprecation_utils3from diffusers.models import cross_attention4utils.deprecate = lambda *arg, **kwargs: None5deprecation_utils.deprecate = lambda *arg, **kwargs: None6cross_attention.deprecate = lambda *arg, **kwargs: None7 8'''9import os10import sys11MAIN_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))12sys.path.insert(0, MAIN_DIR)13os.chdir(MAIN_DIR)14'''15 16import cv217import gradio as gr18import numpy as np19import torch20import random21 22from annotator.util import resize_image, HWC323from annotator.openpose import OpenposeDetector24from diffusers.models.unet_2d_condition import UNet2DConditionModel25from diffusers.pipelines import DiffusionPipeline26from diffusers.schedulers import DPMSolverMultistepScheduler27from models import ControlLoRA, ControlLoRACrossAttnProcessor28 29 30apply_openpose = OpenposeDetector()31 32device = 'cuda' if torch.cuda.is_available() else 'cpu'33 34pipeline = DiffusionPipeline.from_pretrained(35    'IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1', safety_checker=None36)37pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)38pipeline = pipeline.to(device)39unet: UNet2DConditionModel = pipeline.unet40 41ckpt_path = "svjack/pose-control-lora-zh"42control_lora = ControlLoRA.from_pretrained(ckpt_path)43control_lora = control_lora.to(device)44 45# load control lora attention processors46lora_attn_procs = {}47lora_layers_list = list([list(layer_list) for layer_list in control_lora.lora_layers])48n_ch = len(unet.config.block_out_channels)49control_ids = [i for i in range(n_ch)]50for name in pipeline.unet.attn_processors.keys():51    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim52    if name.startswith("mid_block"):53        control_id = control_ids[-1]54    elif name.startswith("up_blocks"):55        block_id = int(name[len("up_blocks.")])56        control_id = list(reversed(control_ids))[block_id]57    elif name.startswith("down_blocks"):58        block_id = int(name[len("down_blocks.")])59        control_id = control_ids[block_id]60 61    lora_layers = lora_layers_list[control_id]62    if len(lora_layers) != 0:63        lora_layer: ControlLoRACrossAttnProcessor = lora_layers.pop(0)64        lora_attn_procs[name] = lora_layer65 66unet.set_attn_processor(lora_attn_procs)67 68 69def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, sample_steps, scale, seed, eta):70    with torch.no_grad():71        input_image = HWC3(input_image)72        detected_map, _ = apply_openpose(resize_image(input_image, detect_resolution))73        detected_map = HWC3(detected_map)74        img = resize_image(input_image, image_resolution)75        H, W, C = img.shape76 77        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)78 79        control = torch.from_numpy(detected_map[...,::-1].copy().transpose([2,0,1])).float().to(device)[None] / 127.5 - 180        _ = control_lora(control).control_states81 82        if seed == -1:83            seed = random.randint(0, 65535)84 85        # run inference86        generator = torch.Generator(device=device).manual_seed(seed)87        images = []88        for i in range(num_samples):89            _ = control_lora(control).control_states90            image = pipeline(91                prompt + ', ' + a_prompt, negative_prompt=n_prompt,92                num_inference_steps=sample_steps, guidance_scale=scale, eta=eta,93                generator=generator, height=H, width=W).images[0]94            images.append(np.asarray(image))95 96        results = images97    return [detected_map] + results98 99 100block = gr.Blocks().queue()101with block:102    with gr.Row():103        gr.Markdown("## Control Stable Diffusion with Human Pose\n")104        gr.Markdown("This _example_ was **drive** from <br/><b><h4>[https://github.com/svjack/ControlLoRA-Chinese](https://github.com/svjack/ControlLoRA-Chinese)</h4></b>\n")105    with gr.Row():106        with gr.Column():107            input_image = gr.Image(source='upload', type="numpy", value = "war_v1.jpg")108            prompt = gr.Textbox(label="Prompt", value = "麦田守望者")109            run_button = gr.Button(label="Run")110            with gr.Accordion("Advanced options", open=False):111                num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)112                image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=256)113                detect_resolution = gr.Slider(label="OpenPose Resolution", minimum=128, maximum=1024, value=512, step=1)114                sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=30, step=1)115                scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)116                seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)117                eta = gr.Number(label="eta", value=0.0)118                a_prompt = gr.Textbox(label="Added Prompt",119                value='详细的模拟混合媒体拼贴画，帆布质地的当代艺术风格，朋克艺术，逼真主义，感性的身体，表现主义，极简主义。杰作，完美的组成，逼真的美丽的脸')120                n_prompt = gr.Textbox(label="Negative Prompt",121                                      value='低质量，模糊，混乱')122        with gr.Column():123            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')124    ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, sample_steps, scale, seed, eta]125    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])126 127 128block.launch(server_name='0.0.0.0')129 130#### block.launch(server_name='172.16.202.228', share=True)131