Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_inference.py 5.25 KB
Newer Older
mashun1's avatar
controlnet  
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149 
from share import *
import config

import cv2
import einops
import numpy as np
import torch
import random
import matplotlib.pyplot as plt

from PIL import Image
from pytorch_lightning import seed_everything
from annotator.util import resize_image, HWC3
from annotator.canny import CannyDetector
from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler


apply_canny = CannyDetector()

model = create_model('./models/cldm_v15.yaml').cpu()
model.load_state_dict(load_state_dict('./models/control_sd15_canny.pth', location='cuda'))
model = model.cuda()
ddim_sampler = DDIMSampler(model)


def process(input_image="test_imgs/bird.png", 
            prompt="bird", 
            a_prompt="best quality, extremely detailed", 
            n_prompt="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality", 
            num_samples=4, 
            image_resolution=512, 
            ddim_steps=20, 
            guess_mode=False, 
            strength=1, 
            scale=9.0, 
            seed=-1, 
            eta=0, 
            low_threshold=100, 
            high_threshold=200):
    with torch.no_grad():
        img = resize_image(HWC3(input_image), image_resolution)
        H, W, C = img.shape

        detected_map = apply_canny(img, low_threshold, high_threshold)
        detected_map = HWC3(detected_map)

        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
        control = torch.stack([control for _ in range(num_samples)], dim=0)
        control = einops.rearrange(control, 'b h w c -> b c h w').clone()

        if seed == -1:
            seed = random.randint(0, 65535)
        seed_everything(seed)

        if config.save_memory:
            model.low_vram_shift(is_diffusing=False)

        cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
        un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
        shape = (4, H // 8, W // 8)

        if config.save_memory:
            model.low_vram_shift(is_diffusing=True)

        model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13)  # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01
        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
                                                     shape, cond, verbose=False, eta=eta,
                                                     unconditional_guidance_scale=scale,
                                                     unconditional_conditioning=un_cond)

        if config.save_memory:
            model.low_vram_shift(is_diffusing=False)

        x_samples = model.decode_first_stage(samples)
        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)

        results = [x_samples[i] for i in range(num_samples)]
    
    return [255 - detected_map] + results


def arange_imgs(images):
    fig, axes = plt.subplots(nrows=1, ncols=len(images), figsize=(12, 4))
    
    for i, image in enumerate(images):
        axes[i].imshow(image)
        axes[i].axis('off')
        
    plt.tight_layout()
    plt.savefig("test_results/test.png")


if __name__ == "__main__":
    from argparse import ArgumentParser
    
    parser = ArgumentParser()
    
    parser.add_argument("--input_image", type=str, default="test_imgs/bird.png", help="输入图片")
    
    parser.add_argument("--prompt", type=str, default="bird")
    
    parser.add_argument("--positive_prompt", type=str, default="best quality, extremely detailed")
    
    parser.add_argument("--negative_prompt", type=str, default="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality")
    
    parser.add_argument("--num_samples", type=int, default=4)
    
    parser.add_argument("--image_resolution", type=int, default=512)
    
    parser.add_argument("--ddim_steps", type=int, default=20)
    
    parser.add_argument("--guess_mode", default=False, action="store_true")
    
    parser.add_argument("--strength", type=float, default=1)
    
    parser.add_argument("--scale", type=float, default=9)

    parser.add_argument("--seed", type=int, default=-1)
    
    parser.add_argument("--eta", type=float, default=0)

    parser.add_argument("--low_threshold", type=float, default=100)

    parser.add_argument("--high_threshold", type=float, default=200)
    
    args = parser.parse_args()
    
    input_image = np.array(Image.open(args.input_image))
    
    images = process(
            input_image,
            prompt=args.prompt,
            a_prompt=args.positive_prompt,
            n_prompt=args.negative_prompt,
            num_samples=args.num_samples,
            image_resolution=args.image_resolution,
            ddim_steps=args.ddim_steps,
            guess_mode=args.guess_mode,
            strength=args.strength,
            scale=args.scale,
            seed=args.seed,
            eta=args.eta,
            low_threshold=args.low_threshold,
            high_threshold=args.high_threshold
        )
    
    arange_imgs(images)