├── .gitignore
├── LICENSE.txt
├── README.md
├── ReleaseNote.md
├── docs
├── en
│ ├── bmab.md
│ ├── manual.md
│ ├── postprocess.md
│ └── preprocess.md
├── kr
│ ├── api.md
│ ├── bmab.md
│ ├── manual.md
│ ├── postprocess.md
│ └── preprocess.md
└── masking.md
├── filter
├── Put filter file here.txt
├── basic.py
└── vintage.py
├── install.py
├── models
└── Put model file here.txt
├── requirements.txt
├── resources
└── preset
│ ├── 3girls.json
│ ├── Put config file here.txt
│ ├── example.json
│ ├── example2.json
│ ├── hand.json
│ └── hand2.json
├── scripts
└── sd_webui_bmab.py
└── sd_bmab
├── __init__.py
├── base
├── __init__.py
├── cache.py
├── common.py
├── context.py
├── detectorbase.py
├── filter.py
├── maskbase.py
├── process.py
├── processorbase.py
└── sam.py
├── bmab.py
├── compat.py
├── constants.py
├── controlnet.py
├── detectors
├── __init__.py
├── anything.py
├── detector.py
├── face.py
├── hand.py
└── person.py
├── external
├── __init__.py
├── groundingdino
│ └── grdino.py
├── iclight
│ ├── __init__.py
│ ├── bmabiclight.py
│ ├── briarmbg.py
│ ├── iclightbg.py
│ └── iclightnm.py
├── kohyahiresfix
│ └── __init__.py
└── lama
│ ├── __init__.py
│ ├── config.yaml
│ └── saicinpainting
│ ├── __init__.py
│ ├── training
│ ├── __init__.py
│ ├── data
│ │ ├── __init__.py
│ │ └── masks.py
│ ├── losses
│ │ ├── __init__.py
│ │ ├── adversarial.py
│ │ ├── constants.py
│ │ ├── distance_weighting.py
│ │ ├── feature_matching.py
│ │ ├── perceptual.py
│ │ ├── segmentation.py
│ │ └── style_loss.py
│ ├── modules
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── depthwise_sep_conv.py
│ │ ├── fake_fakes.py
│ │ ├── ffc.py
│ │ ├── multidilated_conv.py
│ │ ├── multiscale.py
│ │ ├── pix2pixhd.py
│ │ ├── spatial_transform.py
│ │ └── squeeze_excitation.py
│ ├── trainers
│ │ ├── __init__.py
│ │ ├── base.py
│ │ └── default.py
│ └── visualizers
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── colors.py
│ │ ├── directory.py
│ │ └── noop.py
│ └── utils.py
├── masking
├── __init__.py
├── sam.py
└── sam_hq.py
├── parameters.py
├── pipeline
├── __init__.py
├── internal
│ ├── __init__.py
│ ├── intermediate.py
│ └── internalpipe.py
└── post
│ ├── __init__.py
│ └── mainpipe.py
├── processors
├── __init__.py
├── basic
│ ├── __init__.py
│ ├── blend.py
│ ├── edge.py
│ ├── final.py
│ ├── iclight.py
│ ├── img2imgmasking.py
│ ├── intermediate.py
│ └── preprocessfilter.py
├── controlnet
│ ├── __init__.py
│ ├── ipadapter.py
│ ├── noise.py
│ └── pose.py
├── detailer
│ ├── __init__.py
│ ├── face.py
│ ├── hand.py
│ └── person.py
├── postprocess
│ ├── __init__.py
│ ├── finalfilter.py
│ ├── inpaint.py
│ ├── inpaintlama.py
│ ├── upscaleafterprocess.py
│ ├── upscalebeforeprocess.py
│ └── watermark.py
├── preprocess
│ ├── __init__.py
│ ├── preprocessfilter.py
│ ├── pretraining.py
│ ├── refiner.py
│ ├── resample.py
│ └── resize.py
└── utils
│ ├── __init__.py
│ ├── checkpoint.py
│ ├── filesave.py
│ └── modelswitch.py
├── sd_override
├── __init__.py
├── img2img.py
├── samper.py
├── sd_models.py
└── txt2img.py
├── ui.py
└── util
├── __init__.py
└── installhelper.py
/.gitignore:
--------------------------------------------------------------------------------
1 | **/__pycache__/
2 | cache/
3 |
--------------------------------------------------------------------------------
/ReleaseNote.md:
--------------------------------------------------------------------------------
1 | ### v24.05.12
2 |
3 | * New Feature
4 | * ICLight
5 | * Style 추가했습니다.
6 | * intensive : IC-Light original process
7 | * less intensive : 과도한 적용을 줄임
8 | * normal : 배경을 최대한 살리면서 lighting 적용
9 | * soft : 배경을 최대한 살림
10 | * Face : 얼굴 부분을 강조하도록 검출해서 조명으로 강조
11 | * Person : 사람을 찾아 조명으로 강조
12 | * Cache
13 | * 디렉토리가 존재하지 않으면 생성
14 | * Pretraining
15 | * 입력된 폭, 높이가 아니라 실제 이미지의 크기를 사용하도록 수정
16 |
17 | ### v24.05.12
18 |
19 | * New Feature
20 | * ICLight : https://github.com/lllyasviel/IC-Light
21 | * ICLight 기능을 제어할 수 있도록 통합했습니다.
22 | * 메모리 사용량이 많습니다.
23 | * Enable ICLight before upscale 끄면 1200*1800을 4090으로도 어렵습니다.
24 | * Process 작업전 수행되는 preprocess filter가 추가되었습니다.
25 | * Code Fix
26 | * Pose Default 값을 변경했습니다.
27 | * cache를 통합 제어하도록 변경했습니다.
28 |
29 | ### v24.05.07
30 |
31 | * New Feature
32 | * Installer : groudingdino for CUDA support.
33 | * Resample : fix FakeControlNet
34 | * Refactoring : multiple controlnet support
35 |
36 |
37 | ### v24.05.01
38 |
39 | * New Feature
40 | * ControlNet IpAdapter에서 하위 디렉토리를 만들면 하위 디렉토리 안에서 이미지를 적용할 수 있도록 개선했습니다.
41 | * Random으로 지정하면 전체 이미지에서 선택합니다.
42 | * 저장된 정보를 불러왔을때 Pose, IpAdapter 이미지가 표시됩니다.
43 |
44 | * ### v24.05.01
45 |
46 | * New Feature
47 | * ControlNet IpAdapter 기능을 제공합니다.
48 | * sd-webui-controlnet을 이용하여 제공하며, 랜덤하게 이미지를 적용할 수 있도록 개선했습니다.
49 | * Setting > BMAB에 현재 ControlNet IpAdapter의 model을 적어주셔야 합니다.
50 | * 기본적으로 모델은 default로 적용해두었습니다.
51 |
52 | * Code Fix
53 | * Setting에 Additional Checkpoint Path이 sd-webui 1.8에서 동작하지 않아, 이런 경우 빈칸으로 두면 오류가 발생하지 않도록 수정했습니다.
54 | * Person, Face 의 prompt가 hires.fix가 켜져있다면 hires.fix 것을 사용함.
55 |
56 | ### v24.04.30
57 |
58 | * New Feature
59 | * Person Detailer에 checkpoint, vae, sampler, scheduler 및 steps 추가
60 | * Setting에 Additional Checkpoint Path 디렉토리를 입력하면 두 디렉토리 내용을 병합하여 전체적 적용됨.
61 |
62 | * Code Fix
63 | * Person, Face 의 prompt가 hires.fix가 켜져있다면 hires.fix 것을 사용함.
64 |
65 | ### v24.04.26
66 |
67 | * Code Fix
68 | * Pretraining 에 checkpoint, vae 선택 추가
69 | * checkpoint, vae, filter, pose 등에 대한 fresh 버튼을 삭제하고 하나로 통합
70 | * vintage filter 추가
71 |
72 |
73 | ### v24.04.20
74 |
75 | * New Features
76 | * Pose
77 | * Face only 옵션 추가
78 | * Pose를 선택할 수 있도록 추가
79 |
80 | ### v24.04.17
81 |
82 | * New Features
83 | * Pose
84 | * sd-webui-bmab/pose 디렉토리에 있는 이미지를 이용하여 openpose를 동작한다.
85 | * seed에 맞춰서 임의이 순서로 동작한다.
86 | * Face
87 | * CheckPoint를 선택할 수 있도록 추가
88 | * Pretraining
89 | * Filter를 선택할 수 있도록 추가
90 | * Refiner
91 | * VAE를 선택할 수 있도록 추가
92 | * Noise
93 | * Both, Low res fix, High res fix 추가.
94 |
95 | ### v24.04.16
96 |
97 | * Code Fix
98 | * 1.9.0 지원
99 | * 1.8.0 하위 호환
100 | * CheckPoint 관련 전반 수정
101 | * basicsr 설치 오류 수정
102 |
103 | ### v24.04.05
104 |
105 | * Code Fix
106 | * 1.8.0 설치시에 동작하지 않는 버그를 수정했습니다.
107 | * preprocess filter 분리
108 | * hires.fix - upscale 시에 필터가 정확하게 적용되지 않는 버그 수정
109 | * ControlNet Noise 사용시에 그레이 이미지를 사용하도록 수정
110 | * 그 외 몇가지 코드 리펙토링
111 |
112 | ### v23.11.30.0
113 |
114 | * Code Fix
115 | * Support BMAB DR
116 | * Img2img rollback
117 |
118 |
119 | ### v23.11.28.0
120 |
121 | * New Feature
122 | * Kohya Hires.fix
123 | * Preprocess에서 Kohya Hires.fix기능을 추가하였습니다.
124 | * 이 기능을 사용할때 sd 1.5 기준 1024x1025, 1536x1536, 1024x1536, 1536x1024일 경우가 가장 잘나옵니다.
125 | * 이 기능은 원작자가 SDXL을 위해서 만든 기능입니다. 굳이 sd 1.5에서 사용할 필요는 없습니다.
126 |
127 |
128 | ### v23.11.27.0
129 |
130 | * New Feature
131 | * Stop generating gracefully
132 | * BMAB 프로세스가 완료되면 batch가 남아있더라도 종료하는 기능.
133 | * 이미지 생성 중간에 Interrupt 를 눌러서 종료가 아니라, 이미지 생성이 완료되면 종료된다.
134 | * 'Enable BMAB' 오른쪽에 작게 Stop이 있다.
135 | * FinalFilter
136 | * 최종 이미지에 수정을 가할 수 있도록 필터를 적용할 수 있도록 하는 기능.
137 | * 필터는 구현해서 filter에 넣으면 확인할 수 있다.
138 | * BugFix
139 | * Img2Img와 openpose 사용시에 inpaint area 적용되지 않는 문제 수정.
140 | * 약간의 코드 리펙토링
141 |
142 |
143 |
144 |
--------------------------------------------------------------------------------
/docs/en/manual.md:
--------------------------------------------------------------------------------
1 |
2 | ## Quick Test
3 |
4 | Enable을 체크하고 Config Tab에서 Preset "example"을 선택합니다.
5 |
6 | contrast: 1.2
7 | brightness: 0.9
8 | sharpeness: 1.5
9 |
10 | Edge enhancement 적용
11 | Face Detailing 적용
12 | Resize by person 적용
13 |
14 |
15 |
16 | ## 기본 옵션
17 |
18 | Enabled (VERSION): 기능을 켜고 끌 수 있습니다.
19 |
20 | ### Resize and fill override
21 |
22 | Img2Img를 수행하는 경우 "Resize and fill" 을 선택하게 되면
23 | 통상 좌우, 상하로 늘어나거나 비율이 같다면 그대로 크기만 변경됩니다.
24 |
25 | Enabled 된 상태에서는 항상 이미지가 아래에 위치하고,
26 | 왼쪽, 오른쪽, 윗쪽으로 비율에 맞게 늘어납니다.
27 |
28 | 인물의 윗쪽으로 여백이 없는 경우에 적용하면 효과적입니다.
29 | 너무 크게 늘리게 되면 좋은 결과를 얻기 힘듭니다.
30 | 대략 1.1, 1.2 정도 스케일에서 사용하시길 권장합니다.
31 |
32 |
33 | 
34 | 
35 |
36 |
37 |
38 |
39 |
40 |
41 | # Preprocess
42 |
43 | 본 이미지 변경을 하기 전에 사전 처리 과정을 수행합니다.
44 | 조건에 따라 hires.fix 과정에 개입할 수도 있습니다.
45 |
46 | Preprocess
47 |
48 | # BMAB
49 |
50 | Person, Hand, Face detailer를 수행하거나, 이미지 합성 혹은 노이즈 추가등의 기능을 수행합니다.
51 |
52 | bmab
53 |
54 | # Postprocess
55 |
56 | 이미지 처리가 완료된 이후, 인물의 크기에 따라 배경을 확장하거나 upscale을 할 수 있습니다.
57 |
58 | Postprocess
59 |
--------------------------------------------------------------------------------
/docs/en/postprocess.md:
--------------------------------------------------------------------------------
1 |
2 |
3 | # PostProcessing
4 |
5 | ## Resize
6 |
7 | ### Resize by person
8 |
9 | 그림 속 인물중 가장 신장이 큰 사람의 길이와 그림 높이의 비율이 설정값을 넘어가면 비율을 설정값로 맞추는 기능입니다.
10 | 설정값이 0.90이고 인물의 전체 길이: 그림 높이의 비율이 0.95라고 한다면
11 | 배경을 늘려서 인물의 비율이 0.90이 되도록 합니다.
12 | 배경은 왼쪽, 오른쪽, 위쪽으로 늘어납니다.
13 |
14 |
15 | 이 기능은 2가지 방법을 제공하는데 다음과 같습니다.
16 |
17 | #### Inpaint
18 |
19 | Face Detailing과 같은 방법으로 이미지가 완전히 생성되고 난 후에 주변부를 확장합니다.
20 | 이때 이미 생성된 이미지를 훼손하지 않고 주변부만 확장하기 때문에 직관적으로 확인이 가능합니다.
21 | 가장 빠르고 효과적으로 추천합니다.
22 |
23 | #### Inpaint + lama
24 |
25 | Inpaint와 같은 방법인데 BMAB에서 ControlNet을 불러서 Inpaint+lama를 이용해서 동작합니다.
26 | 이미지가 생성되고나서 디테일링 시작하기 전에 img2img를 이용하여 배경을 확장하여 전체적으로 인물을 작게 만드는 효과가 있습니다.
27 |
28 |
29 |
30 |  |
31 |  |
32 |
33 |
34 |  |
35 |  |
36 |
37 |
38 |  |
39 |  |
40 |
41 |
42 |
43 | 이 두가지 방법은 생성된 이미지를 축소하기만 할뿐 훼손하지 않습니다.
44 | 이것이 Resize intermediate와 다른점입니다.
45 |
46 |
47 | ## Upscaler
48 |
49 | 최종적으로 이미지가 완성되고 난 이후에 이미지를 upscaler를 이용하여 크게 만듭니다.
50 |
51 |
52 | #### Enable upscale at final stage
53 |
54 | 이미지 생성이 완료되고 난 이후에 Upscale을 수행합니다.
55 | 960x540으로 생성하고 hires.fix를 x2로 하면 1920x1080 이미지가 나오는데
56 | 여기서 Upscale을 x2를 하면 4K 이미지가 나오게 됩니다.
57 |
58 | #### Detailing after upscale
59 |
60 | 이 옵션을 설정하면 위에서 언급한 Person, Face, Hand 에 대한 detailing을 upscale 이후에 수행합니다.
61 |
62 | #### Upscale Ratio
63 |
64 | 이미지를 얼마나 upscale할지 결정합니다.
65 |
66 |
67 |
--------------------------------------------------------------------------------
/docs/kr/api.md:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 | # API
5 |
6 | Stable diffusion webui의 API를 이용하여 이미지를 생성하는 경우 아래와 같이 BMAB를 사용하여 API Call을 할 수 있습니다.
7 |
8 |
9 | ```python
10 | import requests
11 | import json
12 | import base64
13 |
14 |
15 | prompt = '''
16 | 1girl
17 | '''
18 | negative_prompt = '(worst quality, low quality:1.4),'
19 |
20 | txt2img = {
21 | 'prompt': prompt,
22 | 'negative_prompt': negative_prompt,
23 | 'steps': 20,
24 | 'width': 512,
25 | 'height': 768,
26 | 'cfg_scale': 7,
27 | 'seed': -1,
28 | 'sampler_index': 'DPM++ SDE Karras',
29 | 'script_name': None,
30 | 'alwayson_scripts': {
31 | 'BMAB': {
32 | 'args': [
33 | {
34 | 'enabled': True,
35 | 'face_detailing_enabled': True,
36 | }
37 | ]
38 | }
39 | }
40 | }
41 |
42 | response = requests.post('http://localhost:7860/sdapi/v1/txt2img', data=json.dumps(txt2img))
43 | print(response)
44 | j = response.json()
45 | b64_image = j['images'][0]
46 |
47 |
48 | with open('test.png', 'wb') as image_file:
49 | image_file.write(base64.b64decode(b64_image))
50 |
51 | ```
52 |
53 | BAMB의 Argument는 저장된 설정 파일과 동일하며,
54 | 이것을 기반으로 모든 설정을 사용할 수 있습니다. 설정에 없는 경우 기본값을 사용합니다.
55 |
56 | 아래는 json 형태의 기본 설정 값입니다.
57 |
58 | ```json
59 | {
60 | "enabled": false,
61 | "preprocess_checkpoint": "Use same checkpoint",
62 | "preprocess_vae": "Use same vae",
63 | "txt2img_noise_multiplier": 1,
64 | "txt2img_extra_noise_multiplier": 0,
65 | "txt2img_filter_hresfix_before_upscale": "None",
66 | "txt2img_filter_hresfix_after_upscale": "None",
67 | "resample_enabled": false,
68 | "module_config": {
69 | "resample_opt": {
70 | "save_image": false,
71 | "hiresfix_enabled": false,
72 | "checkpoint": "Use same checkpoint",
73 | "vae": "Use same vae",
74 | "method": "txt2img-1pass",
75 | "filter": "None",
76 | "prompt": "",
77 | "negative_prompt": "",
78 | "sampler": "Use same sampler",
79 | "upscaler": "BMAB fast",
80 | "steps": 20,
81 | "cfg_scale": 7,
82 | "denoising_strength": 0.75,
83 | "strength": 0.5,
84 | "begin": 0.1,
85 | "end": 0.9
86 | },
87 | "pretraining_opt": {
88 | "hiresfix_enabled": false,
89 | "pretraining_model": "Select Model",
90 | "prompt": "",
91 | "negative_prompt": "",
92 | "sampler": "Use same sampler",
93 | "steps": 20,
94 | "cfg_scale": 7,
95 | "denoising_strength": 0.75,
96 | "dilation": 4,
97 | "box_threshold": 0.35
98 | },
99 | "resize_intermediate_opt": {
100 | "resize_by_person": true,
101 | "method": "stretching",
102 | "alignment": "bottom",
103 | "filter": "None",
104 | "scale": 0.85,
105 | "denoising_strength": 0.75
106 | },
107 | "refiner_opt": {
108 | "checkpoint": "Use same checkpoint",
109 | "keep_checkpoint": true,
110 | "prompt": "",
111 | "negative_prompt": "",
112 | "sampler": "Use same sampler",
113 | "upscaler": "BMAB fast",
114 | "steps": 20,
115 | "cfg_scale": 7,
116 | "denoising_strength": 0.75,
117 | "scale": 1,
118 | "width": 0,
119 | "height": 0
120 | },
121 | "person_detailing_opt": {
122 | "best_quality": false,
123 | "force_1:1": false,
124 | "block_overscaled_image": true,
125 | "auto_upscale": true,
126 | "scale": 4,
127 | "dilation": 3,
128 | "area_ratio": 0.1,
129 | "limit": 1,
130 | "background_color": 1,
131 | "background_blur": 0
132 | },
133 | "person_detailing": {
134 | "denoising_strength": 0.4,
135 | "cfg_scale": 7
136 | },
137 | "face_detailing_opt": {
138 | "best_quality": false,
139 | "sort_by": "Score",
140 | "limit": 1,
141 | "prompt0": "",
142 | "negative_prompt0": "",
143 | "prompt1": "",
144 | "negative_prompt1": "",
145 | "prompt2": "",
146 | "negative_prompt2": "",
147 | "prompt3": "",
148 | "negative_prompt3": "",
149 | "prompt4": "",
150 | "negative_prompt4": "",
151 | "override_parameter": false,
152 | "sampler": "Use same sampler",
153 | "detection_model": "BMAB Face(Normal)",
154 | "dilation": 4,
155 | "box_threshold": 0.35,
156 | "skip_large_face": false,
157 | "large_face_pixels": 0.26
158 | },
159 | "face_detailing": {
160 | "width": 512,
161 | "height": 512,
162 | "cfg_scale": 7,
163 | "steps": 20,
164 | "mask_blur": 4,
165 | "inpaint_full_res": "Only masked",
166 | "inpaint_full_res_padding": 32,
167 | "denoising_strength": 0.4
168 | },
169 | "hand_detailing_opt": {
170 | "block_overscaled_image": true,
171 | "best_quality": false,
172 | "detailing_method": "subframe",
173 | "auto_upscale": true,
174 | "scale": 4,
175 | "box_threshold": 0.3,
176 | "dilation": 0.1,
177 | "additional_parameter": ""
178 | },
179 | "hand_detailing": {
180 | "prompt": "",
181 | "negative_prompt": "",
182 | "denoising_strength": 0.4,
183 | "cfg_scale": 7,
184 | "inpaint_full_res": "Only masked",
185 | "inpaint_full_res_padding": 32
186 | },
187 | "controlnet": {
188 | "enabled": false,
189 | "with_refiner": false,
190 | "noise": false,
191 | "noise_strength": 0.4,
192 | "noise_begin": 0.1,
193 | "noise_end": 0.9
194 | },
195 | "resize_by_person_opt": {
196 | "mode": "Inpaint",
197 | "scale": 0.85,
198 | "denoising_strength": 0.6,
199 | "dilation": 30
200 | }
201 | },
202 | "pretraining_enabled": false,
203 | "edge_flavor_enabled": false,
204 | "edge_low_threadhold": 50,
205 | "edge_high_threadhold": 200,
206 | "edge_strength": 0.5,
207 | "resize_intermediate_enabled": false,
208 | "refiner_enabled": false,
209 | "contrast": 1,
210 | "brightness": 1,
211 | "sharpeness": 1,
212 | "color_saturation": 1,
213 | "color_temperature": 0,
214 | "noise_alpha": 0,
215 | "noise_alpha_final": 0,
216 | "input_image": null,
217 | "blend_enabled": false,
218 | "blend_alpha": 1,
219 | "detect_enabled": false,
220 | "masking_prompt": "",
221 | "person_detailing_enabled": false,
222 | "face_detailing_enabled": false,
223 | "face_detailing_before_hiresfix_enabled": false,
224 | "hand_detailing_enabled": false,
225 | "resize_by_person_enabled": false,
226 | "upscale_enabled": false,
227 | "detailing_after_upscale": true,
228 | "upscaler_name": "None",
229 | "upscale_ratio": 1.5,
230 | "config_file": "test",
231 | "preset": "None"
232 | }
233 | ```
234 |
235 |
--------------------------------------------------------------------------------
/docs/kr/manual.md:
--------------------------------------------------------------------------------
1 |
2 | ## Quick Test
3 |
4 | Enable을 체크하고 Config Tab에서 Preset "example"을 선택합니다.
5 |
6 | contrast: 1.2
7 | brightness: 0.9
8 | sharpeness: 1.5
9 |
10 | Edge enhancement 적용
11 | Face Detailing 적용
12 | Resize by person 적용
13 |
14 |
15 |
16 | ## 기본 옵션
17 |
18 | Enabled (VERSION): 기능을 켜고 끌 수 있습니다.
19 |
20 | ### Resize and fill override
21 |
22 | Img2Img를 수행하는 경우 "Resize and fill" 을 선택하게 되면
23 | 통상 좌우, 상하로 늘어나거나 비율이 같다면 그대로 크기만 변경됩니다.
24 |
25 | Enabled 된 상태에서는 항상 이미지가 아래에 위치하고,
26 | 왼쪽, 오른쪽, 윗쪽으로 비율에 맞게 늘어납니다.
27 |
28 | 인물의 윗쪽으로 여백이 없는 경우에 적용하면 효과적입니다.
29 | 너무 크게 늘리게 되면 좋은 결과를 얻기 힘듭니다.
30 | 대략 1.1, 1.2 정도 스케일에서 사용하시길 권장합니다.
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
40 |
41 | # Preprocess
42 |
43 | 본 이미지 변경을 하기 전에 사전 처리 과정을 수행합니다.
44 | 조건에 따라 hires.fix 과정에 개입할 수도 있습니다.
45 |
46 | Preprocess
47 |
48 | # BMAB
49 |
50 | Person, Hand, Face detailer를 수행하거나, 이미지 합성 혹은 노이즈 추가등의 기능을 수행합니다.
51 |
52 | bmab
53 |
54 | # Postprocess
55 |
56 | 이미지 처리가 완료된 이후, 인물의 크기에 따라 배경을 확장하거나 upscale을 할 수 있습니다.
57 |
58 | Postprocess
59 |
60 |
61 | # API
62 |
63 | stable diffusion webui의 API 기능을 호출할때 BMAB extension을 사용할 수 있습니다.
64 |
65 | API Guide
66 |
--------------------------------------------------------------------------------
/docs/kr/postprocess.md:
--------------------------------------------------------------------------------
1 |
2 |
3 | # PostProcessing
4 |
5 | ## Resize
6 |
7 | ### Resize by person
8 |
9 | 그림 속 인물중 가장 신장이 큰 사람의 길이와 그림 높이의 비율이 설정값을 넘어가면 비율을 설정값로 맞추는 기능입니다.
10 | 설정값이 0.90이고 인물의 전체 길이: 그림 높이의 비율이 0.95라고 한다면
11 | 배경을 늘려서 인물의 비율이 0.90이 되도록 합니다.
12 | 배경은 왼쪽, 오른쪽, 위쪽으로 늘어납니다.
13 |
14 |
15 | 이 기능은 2가지 방법을 제공하는데 다음과 같습니다.
16 |
17 | #### Inpaint
18 |
19 | Face Detailing과 같은 방법으로 이미지가 완전히 생성되고 난 후에 주변부를 확장합니다.
20 | 이때 이미 생성된 이미지를 훼손하지 않고 주변부만 확장하기 때문에 직관적으로 확인이 가능합니다.
21 | 가장 빠르고 효과적으로 추천합니다.
22 |
23 | #### Inpaint + lama
24 |
25 | Inpaint와 같은 방법인데 BMAB에서 ControlNet을 불러서 Inpaint+lama를 이용해서 동작합니다.
26 | 이미지가 생성되고나서 디테일링 시작하기 전에 img2img를 이용하여 배경을 확장하여 전체적으로 인물을 작게 만드는 효과가 있습니다.
27 |
28 |
29 |
30 | |
31 | |
32 |
33 |
34 | |
35 | |
36 |
37 |
38 | |
39 | |
40 |
41 |
42 |
43 | 이 두가지 방법은 생성된 이미지를 축소하기만 할뿐 훼손하지 않습니다.
44 | 이것이 Resize intermediate와 다른점입니다.
45 |
46 |
47 | ## Upscaler
48 |
49 | 최종적으로 이미지가 완성되고 난 이후에 이미지를 upscaler를 이용하여 크게 만듭니다.
50 |
51 |
52 | #### Enable upscale at final stage
53 |
54 | 이미지 생성이 완료되고 난 이후에 Upscale을 수행합니다.
55 | 960x540으로 생성하고 hires.fix를 x2로 하면 1920x1080 이미지가 나오는데
56 | 여기서 Upscale을 x2를 하면 4K 이미지가 나오게 됩니다.
57 |
58 | #### Detailing after upscale
59 |
60 | 이 옵션을 설정하면 위에서 언급한 Person, Face, Hand 에 대한 detailing을 upscale 이후에 수행합니다.
61 |
62 | #### Upscale Ratio
63 |
64 | 이미지를 얼마나 upscale할지 결정합니다.
65 |
66 |
67 |
--------------------------------------------------------------------------------
/docs/masking.md:
--------------------------------------------------------------------------------
1 |
2 | # Auto Masking
3 |
4 | ## Detect
5 |
6 | * Img2Img 인물 사진을 추가한다.
7 | * Just resize 선택
8 | * Resize를 하면서 크기를 키워야 좋은 결과를 얻을 수 있다.
9 | 예제의 경우는 512x768 --> 800x1200 으로 resize한 것이다.
10 | * Denoising Strength는 0.4~0.6 정도를 설정한다.
11 | 모델별로 다르지만, 배경까지 심하게 바뀌는 경우 숫자를 줄인다.
12 |
13 |
14 | 
15 |
16 |
17 | * BMAB를 Enable 한다.
18 | * 반드시 Process before img2img를 체크한다.
19 | * Detect enabled 체크하고 prompt에 person, 1girl, human...뭐든..
20 |
21 |
22 | 
23 |
24 |
25 | * 최종적으로 좋은 결과를 내기 위해서 합쳐진 이미지를 프롬프트로 잘 표현해야 한다.
26 | * 프롬프트가 없거나 적당히 적으면 원하는 이미지에서 멀어지기 때문에 프롬프트가 매우 중요하다.
27 |
28 |
29 | 
30 |
31 |
32 |
--------------------------------------------------------------------------------
/filter/Put filter file here.txt:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/filter/Put filter file here.txt
--------------------------------------------------------------------------------
/filter/basic.py:
--------------------------------------------------------------------------------
1 | import os
2 | from PIL import Image
3 | from PIL import ImageEnhance
4 |
5 | import sd_bmab
6 | from sd_bmab import util
7 | from sd_bmab.base import filter
8 | from sd_bmab.base import cache
9 |
10 |
11 | class Filter(filter.BaseFilter):
12 |
13 | def preprocess(self, context, image, *args, **kwargs):
14 | pass
15 |
16 | def basic_process(self, image: Image):
17 | enhancer = ImageEnhance.Brightness(image)
18 | image = enhancer.enhance(0.8)
19 | enhancer = ImageEnhance.Contrast(image)
20 | image = enhancer.enhance(1.2)
21 | return image
22 |
23 | def basic_process_with_noise(self, processed: Image):
24 | noise = cache.get_noise_from_cache(0, processed.width, processed.height).convert('LA')
25 | noise = noise.convert('RGBA')
26 | blended = Image.blend(processed.convert('RGBA'), noise, alpha=0.1)
27 | return self.basic_process(blended.convert('RGB'))
28 |
29 | def process(self, context, image: Image, processed: Image, *args, **kwargs):
30 | print('-----FILTER BASIC-----')
31 | return self.basic_process(processed)
32 |
33 | def postprocess(self, context, *args, **kwargs):
34 | pass
35 |
--------------------------------------------------------------------------------
/filter/vintage.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from PIL import ImageEnhance
3 |
4 | from sd_bmab.base import filter
5 | from sd_bmab.base import cache
6 | from sd_bmab.processors.basic import final
7 |
8 |
9 | CONTRAST = 0.8
10 | BRIGHTNESS = 0.9
11 | SHARPNESS = 0.5
12 | COLOR = 0.85
13 | COLOR_TEMPERATURE = 5240
14 | NOISE = 0.05
15 |
16 |
17 | class Filter(filter.BaseFilter):
18 |
19 | def preprocess(self, context, image, *args, **kwargs):
20 | pass
21 |
22 | def basic_process(self, image: Image):
23 | enhancer = ImageEnhance.Contrast(image)
24 | image = enhancer.enhance(CONTRAST)
25 | enhancer = ImageEnhance.Brightness(image)
26 | image = enhancer.enhance(BRIGHTNESS)
27 | enhancer = ImageEnhance.Sharpness(image)
28 | image = enhancer.enhance(SHARPNESS)
29 | enhancer = ImageEnhance.Color(image)
30 | image = enhancer.enhance(COLOR)
31 | temp = final.calc_color_temperature(COLOR_TEMPERATURE)
32 | az = []
33 | for d in image.getdata():
34 | az.append((int(d[0] * temp[0]), int(d[1] * temp[1]), int(d[2] * temp[2])))
35 | image = Image.new('RGB', image.size)
36 | image.putdata(az)
37 | noise = cache.get_noise_from_cache(0, image.size[0], image.size[1])
38 | image = Image.blend(image, noise, alpha=NOISE)
39 | return image
40 |
41 | def process(self, context, image: Image, processed: Image, *args, **kwargs):
42 | print('-----FILTER VINTAGE-----')
43 | return self.basic_process(processed)
44 |
45 | def postprocess(self, context, *args, **kwargs):
46 | pass
47 |
--------------------------------------------------------------------------------
/install.py:
--------------------------------------------------------------------------------
1 | import json
2 |
3 | import launch
4 |
5 |
6 | def install_segmentanything():
7 | launch.run_pip('install segment_anything')
8 |
9 |
10 | def install_segmentanything_hq():
11 | launch.run_pip('install segment_anything_hq')
12 |
13 |
14 | def install_ultralytics():
15 | launch.run_pip('install ultralytics')
16 |
17 |
18 | def install_diffusers():
19 | launch.run_pip('install diffusers==0.27.2')
20 |
21 |
22 | try:
23 | from basicsr.utils.download_util import load_file_from_url
24 | except:
25 | launch.run_pip('install basicsr==1.4.2')
26 |
27 | required = {
28 | ('segment_anything', install_segmentanything),
29 | ('segment_anything_hq', install_segmentanything_hq),
30 | ('ultralytics', install_ultralytics),
31 | ('diffusers', install_diffusers)
32 | }
33 |
34 | for pack_name, func in required:
35 | if not launch.is_installed(pack_name):
36 | func()
37 |
38 | '''
39 | with open('ui-config.json', 'rt', encoding='UTF8') as f:
40 | j = json.load(f)
41 | print(json.dumps(j, indent=2))
42 | '''
43 |
--------------------------------------------------------------------------------
/models/Put model file here.txt:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/models/Put model file here.txt
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | segment_anything
2 | segment_anything_hq
3 | ultralytics
4 | pillow
5 |
6 | basicsr==1.4.2
7 | kornia
8 | omegaconf
9 | numpy
10 |
--------------------------------------------------------------------------------
/resources/preset/3girls.json:
--------------------------------------------------------------------------------
1 | {
2 | "enabled": true,
3 | "contrast": 1.2,
4 | "brightness": 0.9,
5 | "sharpeness": 1.5,
6 | "execute_before_img2img": true,
7 | "edge_flavor_enabled": true,
8 | "edge_low_threadhold": 50,
9 | "edge_high_threadhold": 200,
10 | "edge_strength": 0.5,
11 | "resize_by_person_enabled": true,
12 | "resize_by_person": 0.85,
13 | "face_detailing_enabled": true,
14 | "module_config": {
15 | "multiple_face": [
16 | {
17 | "denoising_strength": 0.40,
18 | "prompt": "smile, #!org!#",
19 | "inpaint_full_res": true,
20 | "inpaint_full_res_padding": 32,
21 | "cfg_scale": 7
22 | },
23 | {
24 | "denoising_strength": 0.40,
25 | "prompt": "sad, #!org!#",
26 | "inpaint_full_res": true,
27 | "inpaint_full_res_padding": 32,
28 | "cfg_scale": 7
29 | },
30 | {
31 | "denoising_strength": 0.40,
32 | "prompt": "sad, #!org!#",
33 | "width": 512,
34 | "height": 512,
35 | "inpaint_full_res": true,
36 | "inpaint_full_res_padding": 32,
37 | "cfg_scale": 7
38 | }
39 | ],
40 | "multiple_face_opt": {
41 | "mask dilation": 4,
42 | "limit": -1,
43 | "order": "left"
44 | }
45 | }
46 | }
47 |
--------------------------------------------------------------------------------
/resources/preset/Put config file here.txt:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/resources/preset/Put config file here.txt
--------------------------------------------------------------------------------
/resources/preset/example.json:
--------------------------------------------------------------------------------
1 | {
2 | "enabled": true,
3 | "contrast": 1.2,
4 | "brightness": 0.9,
5 | "sharpeness": 1.5,
6 | "edge_flavor_enabled": true,
7 | "edge_low_threadhold": 50,
8 | "edge_high_threadhold": 200,
9 | "edge_strength": 0.5,
10 | "resize_by_person_enabled": true,
11 | "resize_by_person": 0.85,
12 | "face_detailing_enabled": true,
13 | "module_config": {
14 | "face_detailing": {
15 | "denoising_strength": 0.40,
16 | "prompt": "smile, #!org!#",
17 | "inpaint_full_res": true,
18 | "inpaint_full_res_padding": 32,
19 | "cfg_scale": 7
20 | },
21 | "face_detailing_opt": {
22 | "mask dilation": 4
23 | }
24 | }
25 | }
26 |
--------------------------------------------------------------------------------
/resources/preset/example2.json:
--------------------------------------------------------------------------------
1 | {
2 | "enabled": true,
3 | "contrast": 1.2,
4 | "brightness": 0.9,
5 | "sharpeness": 1.5,
6 | "edge_flavor_enabled": true,
7 | "edge_low_threadhold": 50,
8 | "edge_high_threadhold": 200,
9 | "edge_strength": 0.5,
10 | "resize_by_person_enabled": true,
11 | "resize_by_person": 0.85,
12 | "face_detailing_enabled": true,
13 | "module_config": {
14 | "face_detailing": {
15 | "denoising_strength": 0.40,
16 | "prompt": ", (ulzzang-6500:0.4), #!org!#",
17 | "width": 512,
18 | "height": 512,
19 | "inpaint_full_res": true,
20 | "inpaint_full_res_padding": 32,
21 | "cfg_scale": 7
22 | },
23 | "face_detailing_opt": {
24 | "mask dilation": 4
25 | }
26 | }
27 | }
28 |
--------------------------------------------------------------------------------
/resources/preset/hand.json:
--------------------------------------------------------------------------------
1 | {
2 | "enabled": true,
3 | "hand_detailing_enabled": true,
4 | "module_config": {
5 | "hand_detailing": {
6 | "denoising_strength": 0.4,
7 | "steps": 20,
8 | "cfg_scale": 7,
9 | "inpaint_full_res": 0,
10 | "inpaint_full_res_padding": 32
11 | },
12 | "hand_detailing_opt": {
13 | "scale": 2,
14 | "mode": "inpaint",
15 | "detailing method": "subframe"
16 | }
17 | }
18 | }
19 |
--------------------------------------------------------------------------------
/resources/preset/hand2.json:
--------------------------------------------------------------------------------
1 | {
2 | "enabled": true,
3 | "hand_detailing_enabled": true,
4 | "module_config": {
5 | "hand_detailing": {
6 | "denoising_strength": 0.5,
7 | "prompt": "(good anatomy:1.2), (five fingers:1.3), pretty hands, detail hands, detail fingers, detail nails",
8 | "negative_prompt": "(bad anatomy:1.2), (wrong anatomy:1.2), mutation, amputation, extra fingers, missing fingers, disconnected fingers",
9 | "steps": 20,
10 | "cfg_scale": 7,
11 | "inpaint_full_res": 0,
12 | "inpaint_full_res_padding": 32
13 | },
14 | "hand_detailing_opt": {
15 | "scale": 4,
16 | "mode": "inpaint",
17 | "detailing method": "each hand"
18 | }
19 | }
20 | }
21 |
--------------------------------------------------------------------------------
/scripts/sd_webui_bmab.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.bmab import BmabExtScript
2 |
--------------------------------------------------------------------------------
/sd_bmab/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/base/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base.common import VAEMethodOverride
2 | from sd_bmab.base.context import Context
3 | from sd_bmab.base.detectorbase import DetectorBase
4 | from sd_bmab.base.processorbase import ProcessorBase
5 | from sd_bmab.base.process import process_img2img, process_txt2img, build_img2img, apply_extensions, process_img2img_with_controlnet
6 | from sd_bmab.base.maskbase import MaskBase
7 |
8 |
--------------------------------------------------------------------------------
/sd_bmab/base/cache.py:
--------------------------------------------------------------------------------
1 | import os
2 | from PIL import Image
3 |
4 | import sd_bmab
5 | from sd_bmab import util
6 |
7 |
8 | def check_cache_dir(path):
9 | if not os.path.exists(path):
10 | os.mkdir(path)
11 |
12 |
13 | def get_noise_from_cache(seed, width, height):
14 | path = os.path.dirname(sd_bmab.__file__)
15 | path = os.path.normpath(os.path.join(path, '../resources/cache'))
16 | check_cache_dir(path)
17 | cache_file = f'{path}/noise_{width}_{height}.png'
18 | if os.path.isfile(cache_file):
19 | return Image.open(cache_file)
20 | img = util.generate_noise(seed, width, height)
21 | img.save(cache_file)
22 | return img
23 |
24 |
25 | def get_image_from_cache(filename):
26 | path = os.path.dirname(sd_bmab.__file__)
27 | path = os.path.normpath(os.path.join(path, '../resources/cache'))
28 | check_cache_dir(path)
29 | full_path = os.path.join(path, filename)
30 | if os.path.exists(full_path):
31 | return Image.open(full_path)
32 | return None
33 |
34 |
35 | def put_image_to_cache(filename, image):
36 | path = os.path.dirname(sd_bmab.__file__)
37 | path = os.path.normpath(os.path.join(path, '../resources/cache'))
38 | check_cache_dir(path)
39 | full_path = os.path.join(path, filename)
40 | image.save(full_path)
41 |
--------------------------------------------------------------------------------
/sd_bmab/base/common.py:
--------------------------------------------------------------------------------
1 | from modules import shared
2 |
3 |
4 | class VAEMethodOverride:
5 |
6 | def __init__(self, hiresfix=False) -> None:
7 | super().__init__()
8 | self.org_encode_method = None
9 | self.org_decode_method = None
10 | self.img2img_fix_steps = None
11 | self.hiresfix = hiresfix
12 |
13 | def __enter__(self):
14 | if ('sd_vae_encode_method' in shared.opts.data) and shared.opts.bmab_detail_full:
15 | self.encode_method = shared.opts.sd_vae_encode_method
16 | self.decode_method = shared.opts.sd_vae_decode_method
17 | shared.opts.sd_vae_encode_method = 'Full'
18 | shared.opts.sd_vae_decode_method = 'Full'
19 | if self.hiresfix and not shared.opts.img2img_fix_steps:
20 | self.img2img_fix_steps = shared.opts.img2img_fix_steps
21 | shared.opts.img2img_fix_steps = True
22 |
23 | def __exit__(self, *args, **kwargs):
24 | if ('sd_vae_encode_method' in shared.opts.data) and shared.opts.bmab_detail_full:
25 | shared.opts.sd_vae_encode_method = self.encode_method
26 | shared.opts.sd_vae_decode_method = self.decode_method
27 | if self.img2img_fix_steps is not None:
28 | shared.opts.img2img_fix_steps = self.img2img_fix_steps
29 |
30 |
31 | class StopGeneration:
32 |
33 | def __init__(self) -> None:
34 | super().__init__()
35 | if not hasattr(shared.state, 'stopping_generation'):
36 | return
37 | self.stopping_generation = shared.state.stopping_generation
38 |
39 | def __enter__(self):
40 | if not hasattr(shared.state, 'stopping_generation'):
41 | return
42 | shared.state.stopping_generation = False
43 |
44 | def __exit__(self, *args, **kwargs):
45 | if not hasattr(shared.state, 'stopping_generation'):
46 | return
47 | shared.state.stopping_generation = self.stopping_generation
48 |
--------------------------------------------------------------------------------
/sd_bmab/base/context.py:
--------------------------------------------------------------------------------
1 | from modules import shared
2 | from modules.processing import StableDiffusionProcessingImg2Img
3 |
4 | from sd_bmab.sd_override import StableDiffusionProcessingTxt2ImgOv, StableDiffusionProcessingImg2ImgOv
5 | from sd_bmab import constants
6 |
7 |
8 | class Context(object):
9 | def __init__(self, s, p, a, idx, **kwargs) -> None:
10 | super().__init__()
11 |
12 | self.script = s
13 | self.sdprocessing = p
14 | self.args = a
15 | self.index = idx
16 | self.refiner = None
17 | self.sd_model_name = None
18 | self.sd_vae_name = None
19 | self.container = {}
20 |
21 | @staticmethod
22 | def newContext(s, p, a, idx, **kwargs):
23 | return Context(s, p, a, idx, **kwargs)
24 |
25 | def get_current_prompt(self):
26 | return self.sdprocessing.prompt
27 |
28 | def get_prompt_by_index(self):
29 | if self.sdprocessing.all_prompts is None or len(self.sdprocessing.all_prompts) <= self.index:
30 | return self.sdprocessing.prompt
31 | return self.sdprocessing.all_prompts[self.index]
32 |
33 | def get_negative_prompt_by_index(self):
34 | if self.sdprocessing.all_negative_prompts is None or len(self.sdprocessing.all_negative_prompts) <= self.index:
35 | return self.sdprocessing.negative_prompt
36 | return self.sdprocessing.all_negative_prompts[self.index]
37 |
38 | def get_hires_prompt_by_index(self):
39 | if self.sdprocessing.all_hr_prompts is None or len(self.sdprocessing.all_hr_prompts) <= self.index:
40 | return self.sdprocessing.hr_prompt
41 | return self.sdprocessing.all_hr_prompts[self.index]
42 |
43 | def get_hires_negative_prompt_by_index(self):
44 | if self.sdprocessing.all_hr_negative_prompts is None or len(self.sdprocessing.all_hr_negative_prompts) <= self.index:
45 | return self.sdprocessing.hr_negative_prompt
46 | return self.sdprocessing.all_hr_negative_prompts[self.index]
47 |
48 | def get_seeds(self):
49 | if self.sdprocessing.all_seeds is None or self.sdprocessing.all_subseeds is None:
50 | return self.sdprocessing.seed, self.sdprocessing.subseed
51 | if len(self.sdprocessing.all_seeds) <= self.index or len(self.sdprocessing.all_subseeds) <= self.index:
52 | return self.sdprocessing.seed, self.sdprocessing.subseed
53 | return self.sdprocessing.all_seeds[self.index], self.sdprocessing.all_subseeds[self.index]
54 |
55 | def get_max_area(self):
56 | if shared.opts.bmab_optimize_vram == 'low vram':
57 | return 512 * 768
58 | elif shared.opts.bmab_optimize_vram == 'med vram':
59 | return self.sdprocessing.width * self.sdprocessing.height
60 | if isinstance(self.sdprocessing, StableDiffusionProcessingTxt2ImgOv) and self.sdprocessing.enable_hr:
61 | return self.sdprocessing.hr_upscale_to_x * self.sdprocessing.hr_upscale_to_y
62 | return self.sdprocessing.width * self.sdprocessing.height
63 |
64 | def add_generation_param(self, key: object, value: object) -> object:
65 | self.sdprocessing.extra_generation_params[key] = value
66 |
67 | def add_extra_image(self, image):
68 | self.script.extra_image.append(image)
69 |
70 | def with_refiner(self):
71 | return self.args.get('refiner_enabled', False)
72 |
73 | def is_refiner_context(self):
74 | return self.refiner is not None
75 |
76 | def is_hires_fix(self):
77 | if isinstance(self.sdprocessing, StableDiffusionProcessingTxt2ImgOv) and self.sdprocessing.enable_hr:
78 | return True
79 | return False
80 |
81 | def add_job(self, count=1):
82 | shared.state.job_count += count
83 | shared.state.sampling_step = 0
84 | shared.state.current_image_sampling_step = 0
85 |
86 | def is_img2img(self):
87 | return isinstance(self.sdprocessing, StableDiffusionProcessingImg2ImgOv) or isinstance(self.sdprocessing, StableDiffusionProcessingImg2Img)
88 |
89 | def is_txtimg(self):
90 | return isinstance(self.sdprocessing, StableDiffusionProcessingTxt2ImgOv)
91 |
92 | def save_and_apply_checkpoint(self, checkpoint, vae):
93 | self.sd_model_name = checkpoint
94 | self.sd_vae_name = vae
95 |
96 | def restore_checkpoint(self):
97 | self.sd_model_name = None
98 | self.sd_vae_name = None
99 |
100 | def apply_checkpoint(self, options):
101 | if self.sd_model_name is not None and self.sd_model_name != constants.checkpoint_default:
102 | override_settings = options.get('override_settings', {})
103 | override_settings['sd_model_checkpoint'] = self.sd_model_name
104 | options['override_settings'] = override_settings
105 | if self.sd_vae_name is not None and self.sd_vae_name != constants.vae_default:
106 | override_settings = options.get('override_settings', {})
107 | override_settings['sd_vae'] = self.sd_vae_name
108 | options['override_settings'] = override_settings
109 |
110 | def save(self, key, value):
111 | self.container[key] = value
112 |
113 | def load(self, key):
114 | return self.container.get(key)
115 |
--------------------------------------------------------------------------------
/sd_bmab/base/detectorbase.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from sd_bmab.base.context import Context
3 |
4 |
5 | class DetectorBase(object):
6 | def __init__(self, **kwargs) -> None:
7 | super().__init__()
8 |
9 | def target(self):
10 | pass
11 |
12 | def description(self):
13 | pass
14 |
15 | def predict(self, context: Context, image: Image):
16 | pass
17 |
--------------------------------------------------------------------------------
/sd_bmab/base/filter.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import glob
4 | import importlib.util
5 |
6 | from PIL import Image
7 |
8 | import sd_bmab
9 | from sd_bmab import constants
10 | from sd_bmab.util import debug_print
11 | from sd_bmab import controlnet
12 |
13 |
14 | filters = [constants.filter_default]
15 |
16 |
17 | class BaseFilter(object):
18 |
19 | def __init__(self) -> None:
20 | super().__init__()
21 |
22 | def configurations(self):
23 | return {}
24 |
25 | def is_controlnet_required(self):
26 | return False
27 |
28 | def preprocess(self, context, image, *args, **kwargs):
29 | pass
30 |
31 | def process(self, context, base: Image, processed: Image, *args, **kwargs):
32 | return processed
33 |
34 | def postprocess(self, context, *args, **kwargs):
35 | pass
36 |
37 | def finalprocess(self, context, *args, **kwargs):
38 | pass
39 |
40 |
41 | class NoneFilter(BaseFilter):
42 |
43 | def process_filter(self, context, base: Image, processed: Image, *args, **kwargs):
44 | return processed
45 |
46 |
47 | def reload_filters():
48 | global filters
49 | filters = [constants.filter_default]
50 |
51 | path = os.path.dirname(sd_bmab.__file__)
52 | path = os.path.normpath(os.path.join(path, '../filter'))
53 | files = sorted(glob.glob(f'{path}/*.py'))
54 | for file in files:
55 | fname = os.path.splitext(os.path.basename(file))[0]
56 | filters.append(fname)
57 |
58 |
59 | def get_filter(name):
60 | if name == 'None':
61 | return NoneFilter()
62 | debug_print('Filter', name)
63 | path = os.path.dirname(sd_bmab.__file__)
64 | path = os.path.normpath(os.path.join(path, '../filter'))
65 | filter_path = f'{path}/{name}.py'
66 | mod = load_module(filter_path, 'filter')
67 | return eval(f'mod.Filter()')
68 |
69 |
70 | def load_module(file_name, module_name):
71 | spec = importlib.util.spec_from_file_location(module_name, file_name)
72 | module = importlib.util.module_from_spec(spec)
73 | sys.modules[module_name] = module
74 | spec.loader.exec_module(module)
75 | return module
76 |
77 |
78 | def preprocess_filter(bmab_filter, context, image, *args, **kwargs):
79 | bmab_filter.preprocess(context, image, *args, **kwargs)
80 |
81 |
82 | def process_filter(bmab_filter, context, base: Image, processed: Image, *args, **kwargs):
83 | return bmab_filter.process(context, base, processed, *args, **kwargs)
84 |
85 |
86 | def postprocess_filter(bmab_filter, context, *args, **kwargs):
87 | bmab_filter.postprocess(context, *args, **kwargs)
88 |
89 |
90 | def finalprocess_filter(bmab_filter, context, *args, **kwargs):
91 | bmab_filter.finalprocess(context, *args, **kwargs)
92 |
93 |
--------------------------------------------------------------------------------
/sd_bmab/base/maskbase.py:
--------------------------------------------------------------------------------
1 |
2 | class MaskBase(object):
3 | def __init__(self) -> None:
4 | super().__init__()
5 |
6 | @property
7 | def name(self):
8 | pass
9 |
10 | @classmethod
11 | def init(cls, *args, **kwargs):
12 | pass
13 |
14 | def predict(self, image, box):
15 | pass
16 |
17 | def predict_multiple(self, image, points, labels, boxes=None):
18 | pass
19 |
20 | @classmethod
21 | def release(cls):
22 | pass
23 |
--------------------------------------------------------------------------------
/sd_bmab/base/processorbase.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from sd_bmab.base.context import Context
3 |
4 |
5 | class ProcessorBase(object):
6 | def __init__(self) -> None:
7 | super().__init__()
8 |
9 | def use_controlnet(self, context: Context):
10 | return False
11 |
12 | def preprocess(self, context: Context, image: Image):
13 | pass
14 |
15 | def process(self, context: Context, image: Image):
16 | pass
17 |
18 | def postprocess(self, context: Context, image: Image):
19 | pass
20 |
21 | def finalprocess(self, context: Context, image: Image):
22 | pass
23 |
--------------------------------------------------------------------------------
/sd_bmab/base/sam.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import os
3 | import numpy as np
4 |
5 | import torch
6 |
7 | from PIL import Image
8 | from modules.paths import models_path
9 | from modules.safe import unsafe_torch_load, load
10 | from modules.devices import device, torch_gc
11 |
12 | from segment_anything import SamPredictor
13 | from segment_anything import sam_model_registry
14 |
15 | bmab_model_path = os.path.join(models_path, "bmab")
16 |
17 | sam_model = None
18 |
19 |
20 | def sam_init():
21 | MODEL_TYPE = 'vit_b'
22 |
23 | global sam_model
24 | if not sam_model:
25 | torch.load = unsafe_torch_load
26 | sam_model = sam_model_registry[MODEL_TYPE](checkpoint='%s/sam_vit_b_01ec64.pth' % bmab_model_path)
27 | sam_model.to(device=device)
28 | sam_model.eval()
29 | torch.load = load
30 |
31 | return sam_model
32 |
33 |
34 | def sam_predict(pilimg, boxes):
35 | sam = sam_init()
36 |
37 | mask_predictor = SamPredictor(sam)
38 |
39 | numpy_image = np.array(pilimg)
40 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
41 | mask_predictor.set_image(opencv_image)
42 |
43 | result = Image.new('L', pilimg.size, 0)
44 | for box in boxes:
45 | x1, y1, x2, y2 = box
46 |
47 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
48 | masks, scores, logits = mask_predictor.predict(
49 | box=box,
50 | multimask_output=False
51 | )
52 |
53 | mask = Image.fromarray(masks[0])
54 | result.paste(mask, mask=mask)
55 |
56 | return result
57 |
58 |
59 | def sam_predict_box(pilimg, box):
60 | sam = sam_init()
61 |
62 | mask_predictor = SamPredictor(sam)
63 |
64 | numpy_image = np.array(pilimg)
65 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
66 | mask_predictor.set_image(opencv_image)
67 |
68 | x1, y1, x2, y2 = box
69 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
70 |
71 | masks, scores, logits = mask_predictor.predict(
72 | box=box,
73 | multimask_output=False
74 | )
75 |
76 | return Image.fromarray(masks[0])
77 |
78 |
79 | def release():
80 | global sam_model
81 | sam_model = None
82 | torch_gc()
83 |
--------------------------------------------------------------------------------
/sd_bmab/bmab.py:
--------------------------------------------------------------------------------
1 | from modules import scripts
2 | from modules import shared
3 | from modules import script_callbacks
4 | from modules import images
5 |
6 | from sd_bmab import parameters
7 | from sd_bmab.base import context, filter
8 |
9 | from sd_bmab.pipeline import post
10 | from sd_bmab.pipeline import internal
11 | from sd_bmab import masking
12 | from sd_bmab import ui
13 | from sd_bmab import util
14 | from sd_bmab import controlnet
15 | from sd_bmab.sd_override import override_sd_webui, StableDiffusionProcessingTxt2ImgOv
16 | from sd_bmab.sd_override import sd_models
17 | from sd_bmab.compat import check_directory
18 | from sd_bmab.processors.basic import preprocessfilter
19 |
20 |
21 | check_directory()
22 | override_sd_webui()
23 | filter.reload_filters()
24 |
25 | if not shared.opts.data.get('bmab_for_developer', False):
26 | util.check_models()
27 |
28 | if shared.opts.data.get('bmab_additional_checkpoint_path', '') != '':
29 | sd_models.override()
30 |
31 |
32 | class BmabExtScript(scripts.Script):
33 |
34 | def __init__(self) -> None:
35 | super().__init__()
36 | self.extra_image = []
37 | self.config = {}
38 | self.index = 0
39 | self.stop_generation = False
40 |
41 | def title(self):
42 | return 'BMAB'
43 |
44 | def show(self, is_img2img):
45 | return scripts.AlwaysVisible
46 |
47 | def ui(self, is_img2img):
48 | return ui.create_ui(self, is_img2img)
49 |
50 | def before_process(self, p, *args):
51 | self.stop_generation = False
52 | self.extra_image = []
53 | ui.final_images = []
54 | ui.last_process = p
55 | ui.bmab_script = self
56 | self.index = 0
57 | self.config, a = parameters.parse_args(args)
58 | if not a['enabled']:
59 | return
60 |
61 | controlnet.update_controlnet_args(p)
62 | if not hasattr(p, 'context') or p.context is None:
63 | ctx = context.Context.newContext(self, p, a, 0, hiresfix=True)
64 | p.context = ctx
65 | preprocessfilter.run_preprocess_filter(ctx)
66 | post.process_controlnet(p.context)
67 | internal.process_img2img(p.context)
68 | if isinstance(p, StableDiffusionProcessingTxt2ImgOv):
69 | p.initial_noise_multiplier = a.get('txt2img_noise_multiplier', 1)
70 | p.extra_noise = a.get('txt2img_extra_noise_multiplier', 0)
71 |
72 | def postprocess_image(self, p, pp, *args):
73 | self.config, a = parameters.parse_args(args)
74 | if not a['enabled']:
75 | ui.final_images.append(pp.image)
76 | return
77 |
78 | if shared.state.interrupted or shared.state.skipped:
79 | return
80 |
81 | if hasattr(p, 'context') and p.context is not None:
82 | p.context.index = self.index
83 | with controlnet.PreventControlNet(p.context, cn_enabled=post.is_controlnet_required(p.context)):
84 | pp.image = post.process(p.context, pp.image)
85 | ui.final_images.append(pp.image)
86 | self.index += 1
87 | if self.stop_generation:
88 | shared.state.interrupted = True
89 |
90 | def postprocess(self, p, processed, *args):
91 | if shared.opts.bmab_show_extends:
92 | processed.images.extend(self.extra_image)
93 |
94 | post.release()
95 | masking.release()
96 |
97 | def describe(self):
98 | return 'This stuff is worth it, you can buy me a beer in return.'
99 |
100 | def resize_image(self, ctx: context.Context, resize_mode, idx, image, width, height, upscaler_name):
101 | if not ctx.args['enabled']:
102 | return images.resize_image(resize_mode, image, width, height, upscaler_name=upscaler_name)
103 | with controlnet.PreventControlNet(ctx, cn_enabled=internal.is_controlnet_required(ctx)):
104 | image = internal.process_intermediate_before_upscale(ctx, image)
105 | image = images.resize_image(resize_mode, image, width, height, upscaler_name=upscaler_name)
106 | image = internal.process_intermediate_after_upscale(ctx, image)
107 | return image
108 |
109 |
110 | script_callbacks.on_ui_settings(ui.on_ui_settings)
111 |
--------------------------------------------------------------------------------
/sd_bmab/compat.py:
--------------------------------------------------------------------------------
1 | import os
2 | import shutil
3 | import sd_bmab
4 |
5 |
6 | def check_directory():
7 | target = ['cache', 'ipadapter', 'pose', 'saved']
8 | bmab_path = os.path.dirname(sd_bmab.__file__)
9 | dest_path = os.path.normpath(os.path.join(bmab_path, f'../resources'))
10 | for t in target:
11 | path = os.path.normpath(os.path.join(bmab_path, f'../{t}'))
12 | if os.path.exists(path) and os.path.isdir(path):
13 | shutil.move(path, dest_path)
14 |
--------------------------------------------------------------------------------
/sd_bmab/constants.py:
--------------------------------------------------------------------------------
1 | sampler_default = 'Use same sampler'
2 | scheduler_default = 'Use same scheduler'
3 | resize_mode_default = 'Intermediate'
4 | checkpoint_default = 'Use same checkpoint'
5 | vae_default = 'Use same vae'
6 | fast_upscaler = 'BMAB fast'
7 | filter_default = 'None'
8 | face_detector_default = 'BMAB Face(Normal)'
9 |
10 | checkpoint_description = '''
11 |
12 | Specify Checkpoint and VAE to be used in BMAB.
13 | It applies to all functions, and if you change it to Checkpoint and VAE that exist for each function,
14 | it will be applied to all subsequent processes.
15 | '''
16 |
17 |
18 | resize_description = '''
19 |
20 |
21 |
22 | txt2img --resize--> hires.fix --> BMAB Preprocess --> BMAB
23 | txt2img --resize--> BMAB Preprocess --> BMAB
24 |
25 | Methods
26 | stretching : Fast process. Please denoising strength should be over 0.6. (Only CPU).
27 | inpaint : Slow process but High quality. Repaint stretching image. (Use GPU).
28 | inpaint+lama : Very slow process but Very high quality. Repaint stretching image using ControlNet inpaint_only+lama (Use GPU with FIRE!!).
29 | inpaint_only : Very slow process but Very high quality. Repaint stretching image using ControlNet inpaint_only (Use GPU with FIRE!!).
30 |
31 | Please DO NOT SET Latent upscaler in hires.fix.
32 |
33 | '''
34 |
35 | kohya_hiresfix_description = '''
36 |
37 | This is EXPERIMENTAL function.
38 |
39 | '''
40 |
--------------------------------------------------------------------------------
/sd_bmab/controlnet.py:
--------------------------------------------------------------------------------
1 | from copy import copy
2 |
3 | from modules import shared
4 | from modules import processing
5 | from modules import img2img
6 | from modules.processing import Processed
7 |
8 | from sd_bmab.util import debug_print, get_cn_args
9 |
10 |
11 | controlnet_args = (0, 0)
12 |
13 |
14 | class FakeControlNet:
15 | def __init__(self, ctx, cn_enabled=False) -> None:
16 | super().__init__()
17 | self.context = ctx
18 | self.process = ctx.sdprocessing
19 | self.all_prompts = None
20 | self.all_negative_prompts = None
21 | self.extra_image = None
22 | self.enabled = self.is_controlnet_enabled() if cn_enabled else False
23 | self.control_index = []
24 | debug_print('FakeControlNet', self.enabled, cn_enabled)
25 |
26 | def __enter__(self):
27 | if self.enabled:
28 | dummy = Processed(self.process, [], self.process.seed, "")
29 | self.all_prompts = copy(self.process.all_prompts)
30 | self.all_negative_prompts = copy(self.process.all_negative_prompts)
31 | self.extra_image = copy(self.context.script.extra_image)
32 | self.process.scripts.postprocess(copy(self.process), dummy)
33 | for idx, obj in enumerate(self.process.script_args):
34 | if 'controlnet' in obj.__class__.__name__.lower():
35 | if hasattr(obj, 'enabled') and obj.enabled:
36 | obj.enabled = False
37 | self.control_index.append(idx)
38 | elif isinstance(obj, dict) and 'model' in obj and obj['enabled']:
39 | obj['enabled'] = False
40 | self.control_index.append(idx)
41 |
42 | def __exit__(self, *args, **kwargs):
43 | if self.enabled:
44 | copy_p = copy(self.process)
45 | self.process.all_prompts = self.all_prompts
46 | self.process.all_negative_prompts = self.all_negative_prompts
47 | self.extra_image.extend(self.context.script.extra_image)
48 | for idx in self.control_index:
49 | obj = self.process.script_args[idx]
50 | if 'controlnet' in obj.__class__.__name__.lower():
51 | if hasattr(obj, 'enabled'):
52 | obj.enabled = True
53 | elif isinstance(obj, dict) and 'model' in obj:
54 | obj['enabled'] = True
55 | self.process.scripts.before_process(copy_p)
56 | self.process.scripts.process(copy_p)
57 | self.context.script.extra_image.extend(self.extra_image)
58 |
59 | def is_controlnet_enabled(self):
60 | global controlnet_args
61 | for idx in range(controlnet_args[0], controlnet_args[1]):
62 | obj = self.process.script_args[idx]
63 | if isinstance(obj, dict):
64 | return True
65 | if 'controlnet' in obj.__class__.__name__.lower():
66 | if hasattr(obj, 'enabled'):
67 | return True
68 | return False
69 |
70 |
71 | class PreventControlNet(FakeControlNet):
72 | process_images_inner = processing.process_images_inner
73 | process_batch = img2img.process_batch
74 |
75 | def __init__(self, ctx, cn_enabled=False) -> None:
76 | super().__init__(ctx, cn_enabled)
77 | self._process_images_inner = processing.process_images_inner
78 | self._process_batch = img2img.process_batch
79 | self.allow_script_control = None
80 | self.p = ctx.sdprocessing
81 | self.all_prompts = copy(ctx.sdprocessing.all_prompts)
82 | self.all_negative_prompts = copy(ctx.sdprocessing.all_negative_prompts)
83 |
84 | def is_controlnet_used(self):
85 | if not self.p.script_args:
86 | return False
87 |
88 | for idx, obj in enumerate(self.p.script_args):
89 | if 'controlnet' in obj.__class__.__name__.lower():
90 | if hasattr(obj, 'enabled') and obj.enabled:
91 | debug_print('Use controlnet True')
92 | return True
93 | elif isinstance(obj, dict) and 'module' in obj and obj['enabled']:
94 | debug_print('Use controlnet True')
95 | return True
96 |
97 | debug_print('Use controlnet False')
98 | return False
99 |
100 | def __enter__(self):
101 | model = self.p.sd_model.model.diffusion_model
102 | if hasattr(model, '_original_forward'):
103 | model._old_forward = self.p.sd_model.model.diffusion_model.forward
104 | model.forward = getattr(model, '_original_forward')
105 |
106 | processing.process_images_inner = PreventControlNet.process_images_inner
107 | img2img.process_batch = PreventControlNet.process_batch
108 | if 'control_net_allow_script_control' in shared.opts.data:
109 | self.allow_script_control = shared.opts.data['control_net_allow_script_control']
110 | shared.opts.data['control_net_allow_script_control'] = True
111 | self.multiple_tqdm = shared.opts.data.get('multiple_tqdm', True)
112 | shared.opts.data['multiple_tqdm'] = False
113 | super().__enter__()
114 |
115 | def __exit__(self, *args, **kwargs):
116 | processing.process_images_inner = self._process_images_inner
117 | img2img.process_batch = self._process_batch
118 | if 'control_net_allow_script_control' in shared.opts.data:
119 | shared.opts.data['control_net_allow_script_control'] = self.allow_script_control
120 | shared.opts.data['multiple_tqdm'] = self.multiple_tqdm
121 | model = self.p.sd_model.model.diffusion_model
122 | if hasattr(model, '_original_forward') and hasattr(model, '_old_forward'):
123 | self.p.sd_model.model.diffusion_model.forward = model._old_forward
124 | super().__exit__(*args, **kwargs)
125 |
126 |
127 | def update_controlnet_args(p):
128 | cn_arg_index = []
129 | for idx, obj in enumerate(p.script_args):
130 | if 'controlnet' in obj.__class__.__name__.lower():
131 | cn_arg_index.append(idx)
132 | global controlnet_args
133 | controlnet_args = (cn_arg_index[0], cn_arg_index[-1])
134 |
135 |
136 | def get_controlnet_index(p):
137 | cn_args = get_cn_args(p)
138 | controlnet_count = 0
139 | for num in range(*cn_args):
140 | obj = p.script_args[num]
141 | if hasattr(obj, 'enabled') and obj.enabled:
142 | controlnet_count += 1
143 | elif isinstance(obj, dict) and 'model' in obj and obj['enabled']:
144 | controlnet_count += 1
145 | else:
146 | break
147 | return cn_args[0] + controlnet_count
148 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.detectors.face import UltralyticsFaceDetector8s
2 | from sd_bmab.detectors.face import UltralyticsFaceDetector8nv2
3 | from sd_bmab.detectors.face import UltralyticsFaceDetector8n
4 | from sd_bmab.detectors.face import UltralyticsFaceDetector8m
5 | from sd_bmab.detectors.face import BmabFaceNormal
6 | from sd_bmab.detectors.face import BmabFaceSmall
7 | from sd_bmab.detectors.person import UltralyticsPersonDetector8m
8 | from sd_bmab.detectors.person import UltralyticsPersonDetector8n
9 | from sd_bmab.detectors.person import UltralyticsPersonDetector8s
10 | from sd_bmab.detectors.hand import UltralyticsHandDetector8s
11 | from sd_bmab.detectors.hand import UltralyticsHandDetector8n
12 | from sd_bmab.detectors.hand import BmabHandDetector
13 | from sd_bmab.detectors.detector import list_face_detectors, list_hand_detectors, list_person_detectors, get_detector
14 |
15 |
16 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/anything.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base.detectorbase import DetectorBase
2 |
3 |
4 | class UltralyticsAnythingDetector(DetectorBase):
5 | def __init__(self) -> None:
6 | super().__init__()
7 |
8 | def target(self):
9 | pass
10 |
11 | def description(self):
12 | pass
13 |
14 | def detect(self, context, image):
15 | pass
16 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/detector.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base.context import Context
2 |
3 | from sd_bmab.detectors.person import UltralyticsPersonDetector8m
4 | from sd_bmab.detectors.person import UltralyticsPersonDetector8n, UltralyticsPersonDetector8s
5 | from sd_bmab.detectors.face import UltralyticsFaceDetector8n, UltralyticsFaceDetector8s
6 | from sd_bmab.detectors.face import UltralyticsFaceDetector8nv2, UltralyticsFaceDetector8m
7 | from sd_bmab.detectors.face import BmabFaceSmall, BmabFaceNormal
8 | from sd_bmab.detectors.hand import UltralyticsHandDetector8n, UltralyticsHandDetector8s
9 | from sd_bmab.util import debug_print
10 |
11 |
12 | def get_detector(context: Context, model: str, **kwargs):
13 |
14 | debug_print('model', model)
15 | if model == 'face_yolov8n.pt':
16 | return UltralyticsFaceDetector8n(**kwargs)
17 |
18 | all_detectors = [
19 | BmabFaceNormal(**kwargs),
20 | BmabFaceSmall(**kwargs),
21 | UltralyticsPersonDetector8m(**kwargs),
22 | UltralyticsPersonDetector8n(**kwargs),
23 | UltralyticsPersonDetector8s(**kwargs),
24 | UltralyticsFaceDetector8n(**kwargs),
25 | UltralyticsFaceDetector8nv2(**kwargs),
26 | UltralyticsFaceDetector8m(**kwargs),
27 | UltralyticsFaceDetector8s(**kwargs),
28 | UltralyticsHandDetector8n(**kwargs),
29 | UltralyticsHandDetector8s(**kwargs),
30 | ]
31 |
32 | targets = [x for x in all_detectors if model == x.target()]
33 | if len(targets) == 1:
34 | return targets[0]
35 | raise Exception(f'Not found or multiple detector {model}')
36 |
37 |
38 | def list_person_detectors():
39 | kwargs = {}
40 | person_detectors = [
41 | UltralyticsPersonDetector8m(**kwargs),
42 | UltralyticsPersonDetector8n(**kwargs),
43 | UltralyticsPersonDetector8s(**kwargs),
44 | ]
45 | return [x.target() for x in person_detectors]
46 |
47 |
48 | def list_face_detectors():
49 | kwargs = {}
50 | face_detectors = [
51 | BmabFaceNormal(**kwargs),
52 | BmabFaceSmall(**kwargs),
53 | UltralyticsFaceDetector8n(**kwargs),
54 | UltralyticsFaceDetector8nv2(**kwargs),
55 | UltralyticsFaceDetector8m(**kwargs),
56 | UltralyticsFaceDetector8s(**kwargs),
57 | ]
58 | return [x.target() for x in face_detectors]
59 |
60 |
61 | def list_hand_detectors():
62 | kwargs = {}
63 | hand_detectors = [
64 | UltralyticsHandDetector8n(**kwargs),
65 | UltralyticsHandDetector8s(**kwargs),
66 | ]
67 | return [x.target() for x in hand_detectors]
68 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/face.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from PIL import Image
3 | from ultralytics import YOLO
4 |
5 | import modules
6 | from modules import images
7 | from modules import shared
8 |
9 | from sd_bmab import util
10 | from sd_bmab.base.context import Context
11 | from sd_bmab.base.detectorbase import DetectorBase
12 |
13 |
14 | class FaceDetector(DetectorBase):
15 |
16 | def description(self):
17 | return f'Face detecting using {self.target()}'
18 |
19 |
20 | class UltralyticsFaceDetector(FaceDetector):
21 | def __init__(self, **kwargs) -> None:
22 | super().__init__(**kwargs)
23 | self.confidence = kwargs.get('box_threshold', 0.35)
24 | self.model = None
25 |
26 | def target(self):
27 | return f'Ultralytics({self.model})'
28 |
29 | def predict(self, context: Context, image: Image):
30 | yolo = util.lazy_loader(self.model)
31 | boxes = []
32 | confs = []
33 | load = torch.load
34 | torch.load = modules.safe.unsafe_torch_load
35 | try:
36 | model = YOLO(yolo)
37 | pred = model(image, conf=self.confidence, device='')
38 | boxes = pred[0].boxes.xyxy.cpu().numpy()
39 | boxes = boxes.tolist()
40 | confs = pred[0].boxes.conf.tolist()
41 | except:
42 | pass
43 | torch.load = load
44 | return boxes, confs
45 |
46 |
47 | class UltralyticsFaceDetector8n(UltralyticsFaceDetector):
48 | def __init__(self, **kwargs) -> None:
49 | super().__init__(**kwargs)
50 | self.model = 'face_yolov8n.pt'
51 |
52 |
53 | class UltralyticsFaceDetector8m(UltralyticsFaceDetector):
54 | def __init__(self, **kwargs) -> None:
55 | super().__init__(**kwargs)
56 | self.model = 'face_yolov8m.pt'
57 |
58 |
59 | class UltralyticsFaceDetector8nv2(UltralyticsFaceDetector):
60 | def __init__(self, **kwargs) -> None:
61 | super().__init__(**kwargs)
62 | self.model = 'face_yolov8n_v2.pt'
63 |
64 |
65 | class UltralyticsFaceDetector8s(UltralyticsFaceDetector):
66 | def __init__(self, **kwargs) -> None:
67 | super().__init__(**kwargs)
68 | self.model = 'face_yolov8s.pt'
69 |
70 |
71 | class BmabFaceSmall(UltralyticsFaceDetector):
72 | def __init__(self, **kwargs) -> None:
73 | super().__init__(**kwargs)
74 | self.model = 'bmab_face_sm_yolov8n.pt'
75 |
76 | def target(self):
77 | return 'BMAB Face(Small)'
78 |
79 | def predict(self, context: Context, image: Image):
80 | if shared.opts.bmab_debug_logging:
81 | boxes, logits = super().predict(context, image)
82 | if len(boxes) == 0:
83 | images.save_image(
84 | image, context.sdprocessing.outpath_samples, '',
85 | context.sdprocessing.all_seeds[context.index], context.sdprocessing.all_prompts[context.index],
86 | shared.opts.samples_format, p=context.sdprocessing, suffix="-debugging")
87 | det = UltralyticsFaceDetector8n()
88 | return det.predict(context, image)
89 | return boxes, logits
90 | else:
91 | return super().predict(context, image)
92 |
93 |
94 | class BmabFaceNormal(UltralyticsFaceDetector):
95 | def __init__(self, **kwargs) -> None:
96 | super().__init__(**kwargs)
97 | self.model = 'bmab_face_nm_yolov8n.pt'
98 |
99 | def target(self):
100 | return 'BMAB Face(Normal)'
101 |
102 | def predict(self, context: Context, image: Image):
103 | if shared.opts.bmab_debug_logging:
104 | boxes, logits = super().predict(context, image)
105 | if len(boxes) == 0:
106 | images.save_image(
107 | image, context.sdprocessing.outpath_samples, '',
108 | context.sdprocessing.all_seeds[context.index], context.sdprocessing.all_prompts[context.index],
109 | shared.opts.samples_format, p=context.sdprocessing, suffix="-debugging")
110 | det = UltralyticsFaceDetector8n()
111 | return det.predict(context, image)
112 | return boxes, logits
113 | else:
114 | return super().predict(context, image)
115 |
116 |
117 |
118 |
119 |
120 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/hand.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from PIL import Image
3 | from ultralytics import YOLO
4 |
5 | import modules
6 |
7 | from sd_bmab import util
8 | from sd_bmab.base.context import Context
9 | from sd_bmab.base.detectorbase import DetectorBase
10 |
11 |
12 | class HandDetector(DetectorBase):
13 |
14 | def description(self):
15 | return f'Hand detecting using {self.target()}'
16 |
17 |
18 | class UltralyticsHandDetector(HandDetector):
19 | def __init__(self, **kwargs) -> None:
20 | super().__init__(**kwargs)
21 | self.confidence = kwargs.get('box_threshold', 0.35)
22 | self.model = None
23 |
24 | def target(self):
25 | return f'Ultralytics({self.model})'
26 |
27 | def predict(self, context: Context, image: Image):
28 | yolo = util.lazy_loader(self.model)
29 | boxes = []
30 | confs = []
31 | load = torch.load
32 | torch.load = modules.safe.unsafe_torch_load
33 | try:
34 | model = YOLO(yolo)
35 | pred = model(image, conf=self.confidence, device='')
36 | boxes = pred[0].boxes.xyxy.cpu().numpy()
37 | boxes = boxes.tolist()
38 | confs = pred[0].boxes.conf.tolist()
39 | except:
40 | pass
41 | torch.load = load
42 | return boxes, confs
43 |
44 |
45 | class UltralyticsHandDetector8n(UltralyticsHandDetector):
46 | def __init__(self, **kwargs) -> None:
47 | super().__init__(**kwargs)
48 | self.model = 'hand_yolov8n.pt'
49 |
50 |
51 | class UltralyticsHandDetector8s(UltralyticsHandDetector):
52 | def __init__(self, **kwargs) -> None:
53 | super().__init__(**kwargs)
54 | self.model = 'hand_yolov8s.pt'
55 |
56 |
57 | class BmabHandDetector(UltralyticsHandDetector):
58 | def __init__(self, **kwargs) -> None:
59 | super().__init__(**kwargs)
60 | self.model = 'bmab_hand_yolov8n.pt'
61 |
62 | def target(self):
63 | return 'BMAB Hand(Normal)'
64 |
--------------------------------------------------------------------------------
/sd_bmab/detectors/person.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from PIL import Image
3 | from ultralytics import YOLO
4 |
5 | import modules
6 |
7 | from sd_bmab import util
8 | from sd_bmab.base.context import Context
9 | from sd_bmab.base.detectorbase import DetectorBase
10 |
11 |
12 | class PersonDetector(DetectorBase):
13 |
14 | def description(self):
15 | return f'Person detecting using {self.target()}'
16 |
17 |
18 | class UltralyticsPersonDetector(PersonDetector):
19 | def __init__(self, **kwargs) -> None:
20 | super().__init__(**kwargs)
21 | self.confidence = kwargs.get('box_threshold', 0.35)
22 | self.model = None
23 |
24 | def target(self):
25 | return f'Ultralytics({self.model})'
26 |
27 | def predict(self, context: Context, image: Image):
28 | yolo = util.lazy_loader(self.model)
29 | boxes = []
30 | confs = []
31 | load = torch.load
32 | torch.load = modules.safe.unsafe_torch_load
33 | try:
34 | model = YOLO(yolo)
35 | pred = model(image, conf=self.confidence, device='')
36 | boxes = pred[0].boxes.xyxy.cpu().numpy()
37 | boxes = boxes.tolist()
38 | confs = pred[0].boxes.conf.tolist()
39 | except:
40 | pass
41 | torch.load = load
42 | return boxes, confs
43 |
44 |
45 | class UltralyticsPersonDetector8n(UltralyticsPersonDetector):
46 | def __init__(self, **kwargs) -> None:
47 | super().__init__(**kwargs)
48 | self.model = 'person_yolov8n-seg.pt'
49 |
50 |
51 | class UltralyticsPersonDetector8m(UltralyticsPersonDetector):
52 | def __init__(self, **kwargs) -> None:
53 | super().__init__(**kwargs)
54 | self.model = 'person_yolov8m-seg.pt'
55 |
56 |
57 | class UltralyticsPersonDetector8s(UltralyticsPersonDetector):
58 | def __init__(self, **kwargs) -> None:
59 | super().__init__(**kwargs)
60 | self.model = 'person_yolov8s-seg.pt'
61 |
--------------------------------------------------------------------------------
/sd_bmab/external/__init__.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import importlib.util
4 |
5 |
6 | def load_external_module(module, name):
7 | path = os.path.dirname(__file__)
8 | path = os.path.normpath(os.path.join(path, f'{module}/{name}.py'))
9 | return load_module(path, 'module')
10 |
11 |
12 | def load_module(file_name, module_name):
13 | spec = importlib.util.spec_from_file_location(module_name, file_name)
14 | module = importlib.util.module_from_spec(spec)
15 | sys.modules[module_name] = module
16 | spec.loader.exec_module(module)
17 | return module
18 |
19 |
20 | class ModuleAutoLoader(object):
21 |
22 | def __init__(self, module, name) -> None:
23 | super().__init__()
24 | self.mod = load_external_module(module, name)
25 |
26 | def __enter__(self):
27 | return self.mod
28 |
29 | def __exit__(self, *args, **kwargs):
30 | self.mod.release()
31 |
--------------------------------------------------------------------------------
/sd_bmab/external/groundingdino/grdino.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 | import torch
4 |
5 | from groundingdino.util.inference import load_model, predict
6 | from modules.devices import device, torch_gc
7 |
8 | from torchvision.ops import box_convert
9 | import groundingdino.datasets.transforms as T
10 |
11 | from sd_bmab import util
12 |
13 |
14 | dino_model = None
15 |
16 |
17 | def dino_init():
18 | global dino_model
19 | if not dino_model:
20 | util.debug_print('initialize grounding dino.')
21 | swint_ogc = util.lazy_loader('GroundingDINO_SwinT_OGC.py')
22 | swint_ogc_pth = util.lazy_loader('groundingdino_swint_ogc.pth')
23 | dino_model = load_model(swint_ogc, swint_ogc_pth)
24 | return dino_model
25 |
26 |
27 | def dino_predict(pilimg, prompt, box_threahold=0.35, text_threshold=0.25):
28 | transform = T.Compose(
29 | [
30 | T.RandomResize([800], max_size=1333),
31 | T.ToTensor(),
32 | T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
33 | ]
34 | )
35 | img = pilimg.convert('RGB')
36 | image_source = np.asarray(img)
37 | image, _ = transform(img, None)
38 |
39 | model = dino_init()
40 | boxes, logits, phrases = predict(
41 | model=model,
42 | image=image,
43 | caption=prompt,
44 | box_threshold=box_threahold,
45 | text_threshold=text_threshold
46 | )
47 |
48 | h, w, _ = image_source.shape
49 | boxes = boxes * torch.Tensor([w, h, w, h])
50 | annotated_frame = box_convert(boxes=boxes, in_fmt='cxcywh', out_fmt='xyxy').numpy()
51 |
52 | return annotated_frame, logits, phrases
53 |
54 |
55 | def release():
56 | global dino_model
57 | del dino_model
58 | torch_gc()
59 |
--------------------------------------------------------------------------------
/sd_bmab/external/iclight/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/external/iclight/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/external/iclight/bmabiclight.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 |
4 | from ultralytics import YOLO
5 |
6 | from PIL import Image
7 | from PIL import ImageDraw
8 | from PIL import ImageFilter
9 |
10 | import modules
11 | from modules import devices
12 |
13 | from sd_bmab import util
14 | from sd_bmab.base import sam
15 | from sd_bmab.external import load_external_module
16 |
17 |
18 | def process_iclight(context, image, bg_image, prompt, blending, bg_source, arg1, arg2):
19 | np_image = np.array(image.convert('RGB')).astype("uint8")
20 |
21 | if bg_image is None:
22 | mod = load_external_module('iclight', 'iclightnm')
23 | input_fg, matting = mod.run_rmbg(np_image)
24 | seed, subseed = context.get_seeds()
25 | result = mod.process_relight(input_fg, prompt, image.width, image.height, 1, seed, 25,
26 | 'best quality', 'lowres, bad anatomy, bad hands, cropped, worst quality',
27 | arg1[0], arg1[1], arg1[2], arg1[3], bg_source)
28 | mod.release()
29 | fg = Image.fromarray(input_fg.astype('uint8'), 'RGB')
30 | context.add_extra_image(image)
31 | context.add_extra_image(fg)
32 | context.add_extra_image(result)
33 | else:
34 | mod = load_external_module('iclight', 'iclightbg')
35 | input_fg, matting = mod.run_rmbg(np_image)
36 | seed, subseed = context.get_seeds()
37 | result = mod.process_relight(input_fg, None, prompt, image.width, image.height, 1, seed, 20,
38 | 'best quality', 'lowres, bad anatomy, bad hands, cropped, worst quality',
39 | arg2[0], arg2[1], arg2[2], bg_source)
40 | mod.release()
41 | context.add_extra_image(image)
42 | context.add_extra_image(bg_image)
43 | context.add_extra_image(result)
44 | return result
45 |
46 |
47 | def process_bmab_relight(context, image, bg_image, prompt, blending, bg_source, arg1):
48 | mod = load_external_module('iclight', 'iclightbg')
49 | seed, subseed = context.get_seeds()
50 | img1 = image.convert('RGBA')
51 | if bg_image is None:
52 | print('BG Source', bg_source)
53 | if bg_source == 'Face' or bg_source == 'Person':
54 | img2 = generate_detection_gradient(image, bg_source)
55 | context.add_extra_image(img2)
56 | else:
57 | img2 = generate_gradient((32, 32, 32), (224, 224, 224), image.width, image.height, bg_source)
58 | img2 = img2.convert('RGBA')
59 | else:
60 | img2 = bg_image.resize(img1.size, Image.LANCZOS).convert('RGBA')
61 |
62 | blended = Image.blend(img1, img2, alpha=blending)
63 | np_image = np.array(image.convert('RGB')).astype("uint8")
64 | input_bg = np.array(blended.convert('RGB')).astype("uint8")
65 | input_fg, matting = mod.run_rmbg(np_image)
66 | result = mod.process_relight(input_fg, input_bg, prompt, image.width, image.height, 1, seed, 20,
67 | 'best quality', 'lowres, bad anatomy, bad hands, cropped, worst quality',
68 | arg1[0], arg1[1], arg1[2], 'Use Background Image')
69 | mod.release()
70 | return result
71 |
72 |
73 | def generate_gradient(
74 | colour1, colour2, width: int, height: int, d) -> Image:
75 | """Generate a vertical gradient."""
76 | base = Image.new('RGB', (width, height), colour1)
77 | top = Image.new('RGB', (width, height), colour2)
78 | mask = Image.new('L', (width, height))
79 | mask_data = []
80 | if d == 'Left':
81 | for y in range(height):
82 | mask_data.extend([255 - int(255 * (x / width)) for x in range(width)])
83 | if d == 'Right':
84 | for y in range(height):
85 | mask_data.extend([int(255 * (x / width)) for x in range(width)])
86 | if d == 'Bottom':
87 | for y in range(height):
88 | mask_data.extend([int(255 * (y / height))] * width)
89 | if d == 'Top':
90 | for y in range(height):
91 | mask_data.extend([255 - int(255 * (y / height))] * width)
92 | mask.putdata(mask_data)
93 | base.paste(top, (0, 0), mask)
94 | return base
95 |
96 |
97 | def predict(image: Image, model, confidence):
98 | yolo = util.load_pretraining_model(model)
99 | boxes = []
100 | confs = []
101 | load = torch.load
102 | torch.load = modules.safe.unsafe_torch_load
103 | try:
104 | model = YOLO(yolo)
105 | pred = model(image, conf=confidence, device='')
106 | boxes = pred[0].boxes.xyxy.cpu().numpy()
107 | boxes = boxes.tolist()
108 | confs = pred[0].boxes.conf.tolist()
109 | except:
110 | pass
111 | torch.load = load
112 | devices.torch_gc()
113 |
114 | return boxes, confs
115 |
116 |
117 | def generate_detection_gradient(image, model):
118 | mask = Image.new('L', (512, 768), 32)
119 | dr = ImageDraw.Draw(mask, 'L')
120 |
121 | if model == 'Face':
122 | boxes, confs = predict(image, 'face_yolov8n.pt', 0.35)
123 | for box, conf in zip(boxes, confs):
124 | x1, y1, x2, y2 = tuple(int(x) for x in box)
125 | dx = int((x2-x1))
126 | dy = int((y2-y1))
127 | dr.ellipse((x1 - dx, y1 - dy, x2 + dx, y2 + dy), fill=225)
128 | blur = ImageFilter.GaussianBlur(10)
129 | elif model == 'Person':
130 | boxes, confs = predict(image, 'person_yolov8n-seg.pt', 0.35)
131 | for box, conf in zip(boxes, confs):
132 | x1, y1, x2, y2 = tuple(int(x) for x in box)
133 | m = sam.sam_predict_box(image, (x1, y1, x2, y2))
134 | mask.paste(m, mask=m)
135 | blur = ImageFilter.GaussianBlur(30)
136 | else:
137 | return mask
138 | return mask.filter(blur)
139 |
140 |
141 | def bmab_relight(context, process_type, image, bg_image, prompt, blending, bg_source):
142 | if process_type == 'intensive':
143 | if bg_source == 'Face' or bg_source == 'Person':
144 | bg_source = 'None'
145 | return process_iclight(context, image, bg_image, prompt, blending, bg_source, (2, 1.0, 0.5, 0.9), (7, 1.0, 0.5))
146 | elif process_type == 'less intensive':
147 | if bg_source == 'Face' or bg_source == 'Person':
148 | bg_source = 'None'
149 | return process_iclight(context, image, bg_image, prompt, blending, bg_source, (2, 1.0, 0.45, 0.85), (7, 1.0, 0.45))
150 | elif process_type == 'normal':
151 | return process_bmab_relight(context, image, bg_image, prompt, blending, bg_source, (7, 1.0, 0.45))
152 | elif process_type == 'soft':
153 | return process_bmab_relight(context, image, bg_image, prompt, blending, bg_source, (7, 1.0, 0.4))
154 |
155 |
156 | def release():
157 | pass
158 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/__init__.py:
--------------------------------------------------------------------------------
1 | # https://github.com/advimman/lama
2 | # https://github.com/Mikubill/sd-webui-controlnet
3 |
4 | import os
5 | import cv2
6 | import yaml
7 | import torch
8 | import numpy as np
9 | from PIL import Image
10 | from omegaconf import OmegaConf
11 | from einops import rearrange
12 |
13 | from modules import devices
14 |
15 | from sd_bmab.external.lama.saicinpainting.training.trainers import load_checkpoint
16 | from sd_bmab import util
17 |
18 |
19 | class LamaInpainting:
20 |
21 | def __init__(self):
22 | self.model = None
23 | self.device = 'cpu' #devices.get_optimal_device()
24 | # self.device = devices.get_optimal_device()
25 |
26 | @staticmethod
27 | def load_image(pilimg, mode='RGB'):
28 | img = np.array(pilimg.convert(mode))
29 | if img.ndim == 3:
30 | print('transpose')
31 | img = np.transpose(img, (2, 0, 1))
32 | out_img = img.astype('float32') / 255
33 | return out_img
34 |
35 | def load_model(self):
36 | config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.yaml')
37 | cfg = yaml.safe_load(open(config_path, 'rt'))
38 | cfg = OmegaConf.create(cfg)
39 | cfg.training_model.predict_only = True
40 | cfg.visualizer.kind = 'noop'
41 | lamapth = util.lazy_loader('ControlNetLama.pth')
42 | self.model = load_checkpoint(cfg, lamapth, strict=False, map_location='cpu')
43 | self.model = self.model.to(self.device)
44 | self.model.eval()
45 |
46 | def unload_model(self):
47 | if self.model is not None:
48 | self.model.cpu()
49 |
50 | def __call__(self, image, mask):
51 | if self.model is None:
52 | self.load_model()
53 | self.model.to(self.device)
54 |
55 | opencv_image = cv2.cvtColor(np.array(image.convert('RGB')), cv2.COLOR_RGB2BGR)[:, :, 0:3]
56 | opencv_mask = cv2.cvtColor(np.array(mask.convert('RGB')), cv2.COLOR_RGB2BGR)[:, :, 0:1]
57 | color = np.ascontiguousarray(opencv_image).astype(np.float32) / 255.0
58 | mask = np.ascontiguousarray(opencv_mask).astype(np.float32) / 255.0
59 |
60 | with torch.no_grad():
61 | color = torch.from_numpy(color).float().to(self.device)
62 | mask = torch.from_numpy(mask).float().to(self.device)
63 | mask = (mask > 0.5).float()
64 | color = color * (1 - mask)
65 | image_feed = torch.cat([color, mask], dim=2)
66 | image_feed = rearrange(image_feed, 'h w c -> 1 c h w')
67 | result = self.model(image_feed)[0]
68 | result = rearrange(result, 'c h w -> h w c')
69 | result = result * mask + color * (1 - mask)
70 | result *= 255.0
71 |
72 | img = result.detach().cpu().numpy().clip(0, 255).astype(np.uint8)
73 | color_coverted = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
74 | pil_image = Image.fromarray(color_coverted)
75 | return pil_image
76 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/config.yaml:
--------------------------------------------------------------------------------
1 | run_title: b18_ffc075_batch8x15
2 | training_model:
3 | kind: default
4 | visualize_each_iters: 1000
5 | concat_mask: true
6 | store_discr_outputs_for_vis: true
7 | losses:
8 | l1:
9 | weight_missing: 0
10 | weight_known: 10
11 | perceptual:
12 | weight: 0
13 | adversarial:
14 | kind: r1
15 | weight: 10
16 | gp_coef: 0.001
17 | mask_as_fake_target: true
18 | allow_scale_mask: true
19 | feature_matching:
20 | weight: 100
21 | resnet_pl:
22 | weight: 30
23 | weights_path: ${env:TORCH_HOME}
24 |
25 | optimizers:
26 | generator:
27 | kind: adam
28 | lr: 0.001
29 | discriminator:
30 | kind: adam
31 | lr: 0.0001
32 | visualizer:
33 | key_order:
34 | - image
35 | - predicted_image
36 | - discr_output_fake
37 | - discr_output_real
38 | - inpainted
39 | rescale_keys:
40 | - discr_output_fake
41 | - discr_output_real
42 | kind: directory
43 | outdir: /group-volume/User-Driven-Content-Generation/r.suvorov/inpainting/experiments/r.suvorov_2021-04-30_14-41-12_train_simple_pix2pix2_gap_sdpl_novgg_large_b18_ffc075_batch8x15/samples
44 | location:
45 | data_root_dir: /group-volume/User-Driven-Content-Generation/datasets/inpainting_data_root_large
46 | out_root_dir: /group-volume/User-Driven-Content-Generation/${env:USER}/inpainting/experiments
47 | tb_dir: /group-volume/User-Driven-Content-Generation/${env:USER}/inpainting/tb_logs
48 | data:
49 | batch_size: 15
50 | val_batch_size: 2
51 | num_workers: 3
52 | train:
53 | indir: ${location.data_root_dir}/train
54 | out_size: 256
55 | mask_gen_kwargs:
56 | irregular_proba: 1
57 | irregular_kwargs:
58 | max_angle: 4
59 | max_len: 200
60 | max_width: 100
61 | max_times: 5
62 | min_times: 1
63 | box_proba: 1
64 | box_kwargs:
65 | margin: 10
66 | bbox_min_size: 30
67 | bbox_max_size: 150
68 | max_times: 3
69 | min_times: 1
70 | segm_proba: 0
71 | segm_kwargs:
72 | confidence_threshold: 0.5
73 | max_object_area: 0.5
74 | min_mask_area: 0.07
75 | downsample_levels: 6
76 | num_variants_per_mask: 1
77 | rigidness_mode: 1
78 | max_foreground_coverage: 0.3
79 | max_foreground_intersection: 0.7
80 | max_mask_intersection: 0.1
81 | max_hidden_area: 0.1
82 | max_scale_change: 0.25
83 | horizontal_flip: true
84 | max_vertical_shift: 0.2
85 | position_shuffle: true
86 | transform_variant: distortions
87 | dataloader_kwargs:
88 | batch_size: ${data.batch_size}
89 | shuffle: true
90 | num_workers: ${data.num_workers}
91 | val:
92 | indir: ${location.data_root_dir}/val
93 | img_suffix: .png
94 | dataloader_kwargs:
95 | batch_size: ${data.val_batch_size}
96 | shuffle: false
97 | num_workers: ${data.num_workers}
98 | visual_test:
99 | indir: ${location.data_root_dir}/korean_test
100 | img_suffix: _input.png
101 | pad_out_to_modulo: 32
102 | dataloader_kwargs:
103 | batch_size: 1
104 | shuffle: false
105 | num_workers: ${data.num_workers}
106 | generator:
107 | kind: ffc_resnet
108 | input_nc: 4
109 | output_nc: 3
110 | ngf: 64
111 | n_downsampling: 3
112 | n_blocks: 18
113 | add_out_act: sigmoid
114 | init_conv_kwargs:
115 | ratio_gin: 0
116 | ratio_gout: 0
117 | enable_lfu: false
118 | downsample_conv_kwargs:
119 | ratio_gin: ${generator.init_conv_kwargs.ratio_gout}
120 | ratio_gout: ${generator.downsample_conv_kwargs.ratio_gin}
121 | enable_lfu: false
122 | resnet_conv_kwargs:
123 | ratio_gin: 0.75
124 | ratio_gout: ${generator.resnet_conv_kwargs.ratio_gin}
125 | enable_lfu: false
126 | discriminator:
127 | kind: pix2pixhd_nlayer
128 | input_nc: 3
129 | ndf: 64
130 | n_layers: 4
131 | evaluator:
132 | kind: default
133 | inpainted_key: inpainted
134 | integral_kind: ssim_fid100_f1
135 | trainer:
136 | kwargs:
137 | gpus: -1
138 | accelerator: ddp
139 | max_epochs: 200
140 | gradient_clip_val: 1
141 | log_gpu_memory: None
142 | limit_train_batches: 25000
143 | val_check_interval: ${trainer.kwargs.limit_train_batches}
144 | log_every_n_steps: 1000
145 | precision: 32
146 | terminate_on_nan: false
147 | check_val_every_n_epoch: 1
148 | num_sanity_val_steps: 8
149 | limit_val_batches: 1000
150 | replace_sampler_ddp: false
151 | checkpoint_kwargs:
152 | verbose: true
153 | save_top_k: 5
154 | save_last: true
155 | period: 1
156 | monitor: val_ssim_fid100_f1_total_mean
157 | mode: max
158 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/external/lama/saicinpainting/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/external/lama/saicinpainting/training/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/data/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/external/lama/saicinpainting/training/data/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/external/lama/saicinpainting/training/losses/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/constants.py:
--------------------------------------------------------------------------------
1 | weights = {"ade20k":
2 | [6.34517766497462,
3 | 9.328358208955224,
4 | 11.389521640091116,
5 | 16.10305958132045,
6 | 20.833333333333332,
7 | 22.22222222222222,
8 | 25.125628140703515,
9 | 43.29004329004329,
10 | 50.5050505050505,
11 | 54.6448087431694,
12 | 55.24861878453038,
13 | 60.24096385542168,
14 | 62.5,
15 | 66.2251655629139,
16 | 84.74576271186442,
17 | 90.90909090909092,
18 | 91.74311926605505,
19 | 96.15384615384616,
20 | 96.15384615384616,
21 | 97.08737864077669,
22 | 102.04081632653062,
23 | 135.13513513513513,
24 | 149.2537313432836,
25 | 153.84615384615384,
26 | 163.93442622950818,
27 | 166.66666666666666,
28 | 188.67924528301887,
29 | 192.30769230769232,
30 | 217.3913043478261,
31 | 227.27272727272725,
32 | 227.27272727272725,
33 | 227.27272727272725,
34 | 303.03030303030306,
35 | 322.5806451612903,
36 | 333.3333333333333,
37 | 370.3703703703703,
38 | 384.61538461538464,
39 | 416.6666666666667,
40 | 416.6666666666667,
41 | 434.7826086956522,
42 | 434.7826086956522,
43 | 454.5454545454545,
44 | 454.5454545454545,
45 | 500.0,
46 | 526.3157894736842,
47 | 526.3157894736842,
48 | 555.5555555555555,
49 | 555.5555555555555,
50 | 555.5555555555555,
51 | 555.5555555555555,
52 | 555.5555555555555,
53 | 555.5555555555555,
54 | 555.5555555555555,
55 | 588.2352941176471,
56 | 588.2352941176471,
57 | 588.2352941176471,
58 | 588.2352941176471,
59 | 588.2352941176471,
60 | 666.6666666666666,
61 | 666.6666666666666,
62 | 666.6666666666666,
63 | 666.6666666666666,
64 | 714.2857142857143,
65 | 714.2857142857143,
66 | 714.2857142857143,
67 | 714.2857142857143,
68 | 714.2857142857143,
69 | 769.2307692307693,
70 | 769.2307692307693,
71 | 769.2307692307693,
72 | 833.3333333333334,
73 | 833.3333333333334,
74 | 833.3333333333334,
75 | 833.3333333333334,
76 | 909.090909090909,
77 | 1000.0,
78 | 1111.111111111111,
79 | 1111.111111111111,
80 | 1111.111111111111,
81 | 1111.111111111111,
82 | 1111.111111111111,
83 | 1250.0,
84 | 1250.0,
85 | 1250.0,
86 | 1250.0,
87 | 1250.0,
88 | 1428.5714285714287,
89 | 1428.5714285714287,
90 | 1428.5714285714287,
91 | 1428.5714285714287,
92 | 1428.5714285714287,
93 | 1428.5714285714287,
94 | 1428.5714285714287,
95 | 1666.6666666666667,
96 | 1666.6666666666667,
97 | 1666.6666666666667,
98 | 1666.6666666666667,
99 | 1666.6666666666667,
100 | 1666.6666666666667,
101 | 1666.6666666666667,
102 | 1666.6666666666667,
103 | 1666.6666666666667,
104 | 1666.6666666666667,
105 | 1666.6666666666667,
106 | 2000.0,
107 | 2000.0,
108 | 2000.0,
109 | 2000.0,
110 | 2000.0,
111 | 2000.0,
112 | 2000.0,
113 | 2000.0,
114 | 2000.0,
115 | 2000.0,
116 | 2000.0,
117 | 2000.0,
118 | 2000.0,
119 | 2000.0,
120 | 2000.0,
121 | 2000.0,
122 | 2000.0,
123 | 2500.0,
124 | 2500.0,
125 | 2500.0,
126 | 2500.0,
127 | 2500.0,
128 | 2500.0,
129 | 2500.0,
130 | 2500.0,
131 | 2500.0,
132 | 2500.0,
133 | 2500.0,
134 | 2500.0,
135 | 2500.0,
136 | 3333.3333333333335,
137 | 3333.3333333333335,
138 | 3333.3333333333335,
139 | 3333.3333333333335,
140 | 3333.3333333333335,
141 | 3333.3333333333335,
142 | 3333.3333333333335,
143 | 3333.3333333333335,
144 | 3333.3333333333335,
145 | 3333.3333333333335,
146 | 3333.3333333333335,
147 | 3333.3333333333335,
148 | 3333.3333333333335,
149 | 5000.0,
150 | 5000.0,
151 | 5000.0]
152 | }
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/distance_weighting.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | import torchvision
5 |
6 | from sd_bmab.external.lama.saicinpainting.training.losses.perceptual import IMAGENET_STD, IMAGENET_MEAN
7 |
8 |
9 | def dummy_distance_weighter(real_img, pred_img, mask):
10 | return mask
11 |
12 |
13 | def get_gauss_kernel(kernel_size, width_factor=1):
14 | coords = torch.stack(torch.meshgrid(torch.arange(kernel_size),
15 | torch.arange(kernel_size)),
16 | dim=0).float()
17 | diff = torch.exp(-((coords - kernel_size // 2) ** 2).sum(0) / kernel_size / width_factor)
18 | diff /= diff.sum()
19 | return diff
20 |
21 |
22 | class BlurMask(nn.Module):
23 | def __init__(self, kernel_size=5, width_factor=1):
24 | super().__init__()
25 | self.filter = nn.Conv2d(1, 1, kernel_size, padding=kernel_size // 2, padding_mode='replicate', bias=False)
26 | self.filter.weight.data.copy_(get_gauss_kernel(kernel_size, width_factor=width_factor))
27 |
28 | def forward(self, real_img, pred_img, mask):
29 | with torch.no_grad():
30 | result = self.filter(mask) * mask
31 | return result
32 |
33 |
34 | class EmulatedEDTMask(nn.Module):
35 | def __init__(self, dilate_kernel_size=5, blur_kernel_size=5, width_factor=1):
36 | super().__init__()
37 | self.dilate_filter = nn.Conv2d(1, 1, dilate_kernel_size, padding=dilate_kernel_size// 2, padding_mode='replicate',
38 | bias=False)
39 | self.dilate_filter.weight.data.copy_(torch.ones(1, 1, dilate_kernel_size, dilate_kernel_size, dtype=torch.float))
40 | self.blur_filter = nn.Conv2d(1, 1, blur_kernel_size, padding=blur_kernel_size // 2, padding_mode='replicate', bias=False)
41 | self.blur_filter.weight.data.copy_(get_gauss_kernel(blur_kernel_size, width_factor=width_factor))
42 |
43 | def forward(self, real_img, pred_img, mask):
44 | with torch.no_grad():
45 | known_mask = 1 - mask
46 | dilated_known_mask = (self.dilate_filter(known_mask) > 1).float()
47 | result = self.blur_filter(1 - dilated_known_mask) * mask
48 | return result
49 |
50 |
51 | class PropagatePerceptualSim(nn.Module):
52 | def __init__(self, level=2, max_iters=10, temperature=500, erode_mask_size=3):
53 | super().__init__()
54 | vgg = torchvision.models.vgg19(pretrained=True).features
55 | vgg_avg_pooling = []
56 |
57 | for weights in vgg.parameters():
58 | weights.requires_grad = False
59 |
60 | cur_level_i = 0
61 | for module in vgg.modules():
62 | if module.__class__.__name__ == 'Sequential':
63 | continue
64 | elif module.__class__.__name__ == 'MaxPool2d':
65 | vgg_avg_pooling.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0))
66 | else:
67 | vgg_avg_pooling.append(module)
68 | if module.__class__.__name__ == 'ReLU':
69 | cur_level_i += 1
70 | if cur_level_i == level:
71 | break
72 |
73 | self.features = nn.Sequential(*vgg_avg_pooling)
74 |
75 | self.max_iters = max_iters
76 | self.temperature = temperature
77 | self.do_erode = erode_mask_size > 0
78 | if self.do_erode:
79 | self.erode_mask = nn.Conv2d(1, 1, erode_mask_size, padding=erode_mask_size // 2, bias=False)
80 | self.erode_mask.weight.data.fill_(1)
81 |
82 | def forward(self, real_img, pred_img, mask):
83 | with torch.no_grad():
84 | real_img = (real_img - IMAGENET_MEAN.to(real_img)) / IMAGENET_STD.to(real_img)
85 | real_feats = self.features(real_img)
86 |
87 | vertical_sim = torch.exp(-(real_feats[:, :, 1:] - real_feats[:, :, :-1]).pow(2).sum(1, keepdim=True)
88 | / self.temperature)
89 | horizontal_sim = torch.exp(-(real_feats[:, :, :, 1:] - real_feats[:, :, :, :-1]).pow(2).sum(1, keepdim=True)
90 | / self.temperature)
91 |
92 | mask_scaled = F.interpolate(mask, size=real_feats.shape[-2:], mode='bilinear', align_corners=False)
93 | if self.do_erode:
94 | mask_scaled = (self.erode_mask(mask_scaled) > 1).float()
95 |
96 | cur_knowness = 1 - mask_scaled
97 |
98 | for iter_i in range(self.max_iters):
99 | new_top_knowness = F.pad(cur_knowness[:, :, :-1] * vertical_sim, (0, 0, 1, 0), mode='replicate')
100 | new_bottom_knowness = F.pad(cur_knowness[:, :, 1:] * vertical_sim, (0, 0, 0, 1), mode='replicate')
101 |
102 | new_left_knowness = F.pad(cur_knowness[:, :, :, :-1] * horizontal_sim, (1, 0, 0, 0), mode='replicate')
103 | new_right_knowness = F.pad(cur_knowness[:, :, :, 1:] * horizontal_sim, (0, 1, 0, 0), mode='replicate')
104 |
105 | new_knowness = torch.stack([new_top_knowness, new_bottom_knowness,
106 | new_left_knowness, new_right_knowness],
107 | dim=0).max(0).values
108 |
109 | cur_knowness = torch.max(cur_knowness, new_knowness)
110 |
111 | cur_knowness = F.interpolate(cur_knowness, size=mask.shape[-2:], mode='bilinear')
112 | result = torch.min(mask, 1 - cur_knowness)
113 |
114 | return result
115 |
116 |
117 | def make_mask_distance_weighter(kind='none', **kwargs):
118 | if kind == 'none':
119 | return dummy_distance_weighter
120 | if kind == 'blur':
121 | return BlurMask(**kwargs)
122 | if kind == 'edt':
123 | return EmulatedEDTMask(**kwargs)
124 | if kind == 'pps':
125 | return PropagatePerceptualSim(**kwargs)
126 | raise ValueError(f'Unknown mask distance weighter kind {kind}')
127 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/feature_matching.py:
--------------------------------------------------------------------------------
1 | from typing import List
2 |
3 | import torch
4 | import torch.nn.functional as F
5 |
6 |
7 | def masked_l2_loss(pred, target, mask, weight_known, weight_missing):
8 | per_pixel_l2 = F.mse_loss(pred, target, reduction='none')
9 | pixel_weights = mask * weight_missing + (1 - mask) * weight_known
10 | return (pixel_weights * per_pixel_l2).mean()
11 |
12 |
13 | def masked_l1_loss(pred, target, mask, weight_known, weight_missing):
14 | per_pixel_l1 = F.l1_loss(pred, target, reduction='none')
15 | pixel_weights = mask * weight_missing + (1 - mask) * weight_known
16 | return (pixel_weights * per_pixel_l1).mean()
17 |
18 |
19 | def feature_matching_loss(fake_features: List[torch.Tensor], target_features: List[torch.Tensor], mask=None):
20 | if mask is None:
21 | res = torch.stack([F.mse_loss(fake_feat, target_feat)
22 | for fake_feat, target_feat in zip(fake_features, target_features)]).mean()
23 | else:
24 | res = 0
25 | norm = 0
26 | for fake_feat, target_feat in zip(fake_features, target_features):
27 | cur_mask = F.interpolate(mask, size=fake_feat.shape[-2:], mode='bilinear', align_corners=False)
28 | error_weights = 1 - cur_mask
29 | cur_val = ((fake_feat - target_feat).pow(2) * error_weights).mean()
30 | res = res + cur_val
31 | norm += 1
32 | res = res / norm
33 | return res
34 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/perceptual.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | import torchvision
5 |
6 | # from models.ade20k import ModelBuilder
7 | from sd_bmab.external.lama.saicinpainting.utils import check_and_warn_input_range
8 |
9 |
10 | IMAGENET_MEAN = torch.FloatTensor([0.485, 0.456, 0.406])[None, :, None, None]
11 | IMAGENET_STD = torch.FloatTensor([0.229, 0.224, 0.225])[None, :, None, None]
12 |
13 |
14 | class PerceptualLoss(nn.Module):
15 | def __init__(self, normalize_inputs=True):
16 | super(PerceptualLoss, self).__init__()
17 |
18 | self.normalize_inputs = normalize_inputs
19 | self.mean_ = IMAGENET_MEAN
20 | self.std_ = IMAGENET_STD
21 |
22 | vgg = torchvision.models.vgg19(pretrained=True).features
23 | vgg_avg_pooling = []
24 |
25 | for weights in vgg.parameters():
26 | weights.requires_grad = False
27 |
28 | for module in vgg.modules():
29 | if module.__class__.__name__ == 'Sequential':
30 | continue
31 | elif module.__class__.__name__ == 'MaxPool2d':
32 | vgg_avg_pooling.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0))
33 | else:
34 | vgg_avg_pooling.append(module)
35 |
36 | self.vgg = nn.Sequential(*vgg_avg_pooling)
37 |
38 | def do_normalize_inputs(self, x):
39 | return (x - self.mean_.to(x.device)) / self.std_.to(x.device)
40 |
41 | def partial_losses(self, input, target, mask=None):
42 | check_and_warn_input_range(target, 0, 1, 'PerceptualLoss target in partial_losses')
43 |
44 | # we expect input and target to be in [0, 1] range
45 | losses = []
46 |
47 | if self.normalize_inputs:
48 | features_input = self.do_normalize_inputs(input)
49 | features_target = self.do_normalize_inputs(target)
50 | else:
51 | features_input = input
52 | features_target = target
53 |
54 | for layer in self.vgg[:30]:
55 |
56 | features_input = layer(features_input)
57 | features_target = layer(features_target)
58 |
59 | if layer.__class__.__name__ == 'ReLU':
60 | loss = F.mse_loss(features_input, features_target, reduction='none')
61 |
62 | if mask is not None:
63 | cur_mask = F.interpolate(mask, size=features_input.shape[-2:],
64 | mode='bilinear', align_corners=False)
65 | loss = loss * (1 - cur_mask)
66 |
67 | loss = loss.mean(dim=tuple(range(1, len(loss.shape))))
68 | losses.append(loss)
69 |
70 | return losses
71 |
72 | def forward(self, input, target, mask=None):
73 | losses = self.partial_losses(input, target, mask=mask)
74 | return torch.stack(losses).sum(dim=0)
75 |
76 | def get_global_features(self, input):
77 | check_and_warn_input_range(input, 0, 1, 'PerceptualLoss input in get_global_features')
78 |
79 | if self.normalize_inputs:
80 | features_input = self.do_normalize_inputs(input)
81 | else:
82 | features_input = input
83 |
84 | features_input = self.vgg(features_input)
85 | return features_input
86 |
87 |
88 | class ResNetPL(nn.Module):
89 | def __init__(self, weight=1,
90 | weights_path=None, arch_encoder='resnet50dilated', segmentation=True):
91 | super().__init__()
92 | self.impl = ModelBuilder.get_encoder(weights_path=weights_path,
93 | arch_encoder=arch_encoder,
94 | arch_decoder='ppm_deepsup',
95 | fc_dim=2048,
96 | segmentation=segmentation)
97 | self.impl.eval()
98 | for w in self.impl.parameters():
99 | w.requires_grad_(False)
100 |
101 | self.weight = weight
102 |
103 | def forward(self, pred, target):
104 | pred = (pred - IMAGENET_MEAN.to(pred)) / IMAGENET_STD.to(pred)
105 | target = (target - IMAGENET_MEAN.to(target)) / IMAGENET_STD.to(target)
106 |
107 | pred_feats = self.impl(pred, return_feature_maps=True)
108 | target_feats = self.impl(target, return_feature_maps=True)
109 |
110 | result = torch.stack([F.mse_loss(cur_pred, cur_target)
111 | for cur_pred, cur_target
112 | in zip(pred_feats, target_feats)]).sum() * self.weight
113 | return result
114 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/segmentation.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 |
5 | from .constants import weights as constant_weights
6 |
7 |
8 | class CrossEntropy2d(nn.Module):
9 | def __init__(self, reduction="mean", ignore_label=255, weights=None, *args, **kwargs):
10 | """
11 | weight (Tensor, optional): a manual rescaling weight given to each class.
12 | If given, has to be a Tensor of size "nclasses"
13 | """
14 | super(CrossEntropy2d, self).__init__()
15 | self.reduction = reduction
16 | self.ignore_label = ignore_label
17 | self.weights = weights
18 | if self.weights is not None:
19 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
20 | self.weights = torch.FloatTensor(constant_weights[weights]).to(device)
21 |
22 | def forward(self, predict, target):
23 | """
24 | Args:
25 | predict:(n, c, h, w)
26 | target:(n, 1, h, w)
27 | """
28 | target = target.long()
29 | assert not target.requires_grad
30 | assert predict.dim() == 4, "{0}".format(predict.size())
31 | assert target.dim() == 4, "{0}".format(target.size())
32 | assert predict.size(0) == target.size(0), "{0} vs {1} ".format(predict.size(0), target.size(0))
33 | assert target.size(1) == 1, "{0}".format(target.size(1))
34 | assert predict.size(2) == target.size(2), "{0} vs {1} ".format(predict.size(2), target.size(2))
35 | assert predict.size(3) == target.size(3), "{0} vs {1} ".format(predict.size(3), target.size(3))
36 | target = target.squeeze(1)
37 | n, c, h, w = predict.size()
38 | target_mask = (target >= 0) * (target != self.ignore_label)
39 | target = target[target_mask]
40 | predict = predict.transpose(1, 2).transpose(2, 3).contiguous()
41 | predict = predict[target_mask.view(n, h, w, 1).repeat(1, 1, 1, c)].view(-1, c)
42 | loss = F.cross_entropy(predict, target, weight=self.weights, reduction=self.reduction)
43 | return loss
44 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/losses/style_loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torchvision.models as models
4 |
5 |
6 | class PerceptualLoss(nn.Module):
7 | r"""
8 | Perceptual loss, VGG-based
9 | https://arxiv.org/abs/1603.08155
10 | https://github.com/dxyang/StyleTransfer/blob/master/utils.py
11 | """
12 |
13 | def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]):
14 | super(PerceptualLoss, self).__init__()
15 | self.add_module('vgg', VGG19())
16 | self.criterion = torch.nn.L1Loss()
17 | self.weights = weights
18 |
19 | def __call__(self, x, y):
20 | # Compute features
21 | x_vgg, y_vgg = self.vgg(x), self.vgg(y)
22 |
23 | content_loss = 0.0
24 | content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1'])
25 | content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1'])
26 | content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1'])
27 | content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1'])
28 | content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])
29 |
30 |
31 | return content_loss
32 |
33 |
34 | class VGG19(torch.nn.Module):
35 | def __init__(self):
36 | super(VGG19, self).__init__()
37 | features = models.vgg19(pretrained=True).features
38 | self.relu1_1 = torch.nn.Sequential()
39 | self.relu1_2 = torch.nn.Sequential()
40 |
41 | self.relu2_1 = torch.nn.Sequential()
42 | self.relu2_2 = torch.nn.Sequential()
43 |
44 | self.relu3_1 = torch.nn.Sequential()
45 | self.relu3_2 = torch.nn.Sequential()
46 | self.relu3_3 = torch.nn.Sequential()
47 | self.relu3_4 = torch.nn.Sequential()
48 |
49 | self.relu4_1 = torch.nn.Sequential()
50 | self.relu4_2 = torch.nn.Sequential()
51 | self.relu4_3 = torch.nn.Sequential()
52 | self.relu4_4 = torch.nn.Sequential()
53 |
54 | self.relu5_1 = torch.nn.Sequential()
55 | self.relu5_2 = torch.nn.Sequential()
56 | self.relu5_3 = torch.nn.Sequential()
57 | self.relu5_4 = torch.nn.Sequential()
58 |
59 | for x in range(2):
60 | self.relu1_1.add_module(str(x), features[x])
61 |
62 | for x in range(2, 4):
63 | self.relu1_2.add_module(str(x), features[x])
64 |
65 | for x in range(4, 7):
66 | self.relu2_1.add_module(str(x), features[x])
67 |
68 | for x in range(7, 9):
69 | self.relu2_2.add_module(str(x), features[x])
70 |
71 | for x in range(9, 12):
72 | self.relu3_1.add_module(str(x), features[x])
73 |
74 | for x in range(12, 14):
75 | self.relu3_2.add_module(str(x), features[x])
76 |
77 | for x in range(14, 16):
78 | self.relu3_2.add_module(str(x), features[x])
79 |
80 | for x in range(16, 18):
81 | self.relu3_4.add_module(str(x), features[x])
82 |
83 | for x in range(18, 21):
84 | self.relu4_1.add_module(str(x), features[x])
85 |
86 | for x in range(21, 23):
87 | self.relu4_2.add_module(str(x), features[x])
88 |
89 | for x in range(23, 25):
90 | self.relu4_3.add_module(str(x), features[x])
91 |
92 | for x in range(25, 27):
93 | self.relu4_4.add_module(str(x), features[x])
94 |
95 | for x in range(27, 30):
96 | self.relu5_1.add_module(str(x), features[x])
97 |
98 | for x in range(30, 32):
99 | self.relu5_2.add_module(str(x), features[x])
100 |
101 | for x in range(32, 34):
102 | self.relu5_3.add_module(str(x), features[x])
103 |
104 | for x in range(34, 36):
105 | self.relu5_4.add_module(str(x), features[x])
106 |
107 | # don't need the gradients, just want the features
108 | for param in self.parameters():
109 | param.requires_grad = False
110 |
111 | def forward(self, x):
112 | relu1_1 = self.relu1_1(x)
113 | relu1_2 = self.relu1_2(relu1_1)
114 |
115 | relu2_1 = self.relu2_1(relu1_2)
116 | relu2_2 = self.relu2_2(relu2_1)
117 |
118 | relu3_1 = self.relu3_1(relu2_2)
119 | relu3_2 = self.relu3_2(relu3_1)
120 | relu3_3 = self.relu3_3(relu3_2)
121 | relu3_4 = self.relu3_4(relu3_3)
122 |
123 | relu4_1 = self.relu4_1(relu3_4)
124 | relu4_2 = self.relu4_2(relu4_1)
125 | relu4_3 = self.relu4_3(relu4_2)
126 | relu4_4 = self.relu4_4(relu4_3)
127 |
128 | relu5_1 = self.relu5_1(relu4_4)
129 | relu5_2 = self.relu5_2(relu5_1)
130 | relu5_3 = self.relu5_3(relu5_2)
131 | relu5_4 = self.relu5_4(relu5_3)
132 |
133 | out = {
134 | 'relu1_1': relu1_1,
135 | 'relu1_2': relu1_2,
136 |
137 | 'relu2_1': relu2_1,
138 | 'relu2_2': relu2_2,
139 |
140 | 'relu3_1': relu3_1,
141 | 'relu3_2': relu3_2,
142 | 'relu3_3': relu3_3,
143 | 'relu3_4': relu3_4,
144 |
145 | 'relu4_1': relu4_1,
146 | 'relu4_2': relu4_2,
147 | 'relu4_3': relu4_3,
148 | 'relu4_4': relu4_4,
149 |
150 | 'relu5_1': relu5_1,
151 | 'relu5_2': relu5_2,
152 | 'relu5_3': relu5_3,
153 | 'relu5_4': relu5_4,
154 | }
155 | return out
156 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/__init__.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 | from sd_bmab.external.lama.saicinpainting.training.modules.ffc import FFCResNetGenerator
4 | from sd_bmab.external.lama.saicinpainting.training.modules.pix2pixhd import GlobalGenerator, MultiDilatedGlobalGenerator, \
5 | NLayerDiscriminator, MultidilatedNLayerDiscriminator
6 |
7 |
8 | def make_generator(config, kind, **kwargs):
9 | logging.info(f'Make generator {kind}')
10 |
11 | if kind == 'pix2pixhd_multidilated':
12 | return MultiDilatedGlobalGenerator(**kwargs)
13 |
14 | if kind == 'pix2pixhd_global':
15 | return GlobalGenerator(**kwargs)
16 |
17 | if kind == 'ffc_resnet':
18 | return FFCResNetGenerator(**kwargs)
19 |
20 | raise ValueError(f'Unknown generator kind {kind}')
21 |
22 |
23 | def make_discriminator(kind, **kwargs):
24 | logging.info(f'Make discriminator {kind}')
25 |
26 | if kind == 'pix2pixhd_nlayer_multidilated':
27 | return MultidilatedNLayerDiscriminator(**kwargs)
28 |
29 | if kind == 'pix2pixhd_nlayer':
30 | return NLayerDiscriminator(**kwargs)
31 |
32 | raise ValueError(f'Unknown discriminator kind {kind}')
33 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/base.py:
--------------------------------------------------------------------------------
1 | import abc
2 | from typing import Tuple, List
3 |
4 | import torch
5 | import torch.nn as nn
6 |
7 | from sd_bmab.external.lama.saicinpainting.training.modules.depthwise_sep_conv import DepthWiseSeperableConv
8 | from sd_bmab.external.lama.saicinpainting.training.modules.multidilated_conv import MultidilatedConv
9 |
10 |
11 | class BaseDiscriminator(nn.Module):
12 | @abc.abstractmethod
13 | def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
14 | """
15 | Predict scores and get intermediate activations. Useful for feature matching loss
16 | :return tuple (scores, list of intermediate activations)
17 | """
18 | raise NotImplemented()
19 |
20 |
21 | def get_conv_block_ctor(kind='default'):
22 | if not isinstance(kind, str):
23 | return kind
24 | if kind == 'default':
25 | return nn.Conv2d
26 | if kind == 'depthwise':
27 | return DepthWiseSeperableConv
28 | if kind == 'multidilated':
29 | return MultidilatedConv
30 | raise ValueError(f'Unknown convolutional block kind {kind}')
31 |
32 |
33 | def get_norm_layer(kind='bn'):
34 | if not isinstance(kind, str):
35 | return kind
36 | if kind == 'bn':
37 | return nn.BatchNorm2d
38 | if kind == 'in':
39 | return nn.InstanceNorm2d
40 | raise ValueError(f'Unknown norm block kind {kind}')
41 |
42 |
43 | def get_activation(kind='tanh'):
44 | if kind == 'tanh':
45 | return nn.Tanh()
46 | if kind == 'sigmoid':
47 | return nn.Sigmoid()
48 | if kind is False:
49 | return nn.Identity()
50 | raise ValueError(f'Unknown activation kind {kind}')
51 |
52 |
53 | class SimpleMultiStepGenerator(nn.Module):
54 | def __init__(self, steps: List[nn.Module]):
55 | super().__init__()
56 | self.steps = nn.ModuleList(steps)
57 |
58 | def forward(self, x):
59 | cur_in = x
60 | outs = []
61 | for step in self.steps:
62 | cur_out = step(cur_in)
63 | outs.append(cur_out)
64 | cur_in = torch.cat((cur_in, cur_out), dim=1)
65 | return torch.cat(outs[::-1], dim=1)
66 |
67 | def deconv_factory(kind, ngf, mult, norm_layer, activation, max_features):
68 | if kind == 'convtranspose':
69 | return [nn.ConvTranspose2d(min(max_features, ngf * mult),
70 | min(max_features, int(ngf * mult / 2)),
71 | kernel_size=3, stride=2, padding=1, output_padding=1),
72 | norm_layer(min(max_features, int(ngf * mult / 2))), activation]
73 | elif kind == 'bilinear':
74 | return [nn.Upsample(scale_factor=2, mode='bilinear'),
75 | DepthWiseSeperableConv(min(max_features, ngf * mult),
76 | min(max_features, int(ngf * mult / 2)),
77 | kernel_size=3, stride=1, padding=1),
78 | norm_layer(min(max_features, int(ngf * mult / 2))), activation]
79 | else:
80 | raise Exception(f"Invalid deconv kind: {kind}")
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/depthwise_sep_conv.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | class DepthWiseSeperableConv(nn.Module):
5 | def __init__(self, in_dim, out_dim, *args, **kwargs):
6 | super().__init__()
7 | if 'groups' in kwargs:
8 | # ignoring groups for Depthwise Sep Conv
9 | del kwargs['groups']
10 |
11 | self.depthwise = nn.Conv2d(in_dim, in_dim, *args, groups=in_dim, **kwargs)
12 | self.pointwise = nn.Conv2d(in_dim, out_dim, kernel_size=1)
13 |
14 | def forward(self, x):
15 | out = self.depthwise(x)
16 | out = self.pointwise(out)
17 | return out
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/fake_fakes.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from kornia import SamplePadding
3 | from kornia.augmentation import RandomAffine, CenterCrop
4 |
5 |
6 | class FakeFakesGenerator:
7 | def __init__(self, aug_proba=0.5, img_aug_degree=30, img_aug_translate=0.2):
8 | self.grad_aug = RandomAffine(degrees=360,
9 | translate=0.2,
10 | padding_mode=SamplePadding.REFLECTION,
11 | keepdim=False,
12 | p=1)
13 | self.img_aug = RandomAffine(degrees=img_aug_degree,
14 | translate=img_aug_translate,
15 | padding_mode=SamplePadding.REFLECTION,
16 | keepdim=True,
17 | p=1)
18 | self.aug_proba = aug_proba
19 |
20 | def __call__(self, input_images, masks):
21 | blend_masks = self._fill_masks_with_gradient(masks)
22 | blend_target = self._make_blend_target(input_images)
23 | result = input_images * (1 - blend_masks) + blend_target * blend_masks
24 | return result, blend_masks
25 |
26 | def _make_blend_target(self, input_images):
27 | batch_size = input_images.shape[0]
28 | permuted = input_images[torch.randperm(batch_size)]
29 | augmented = self.img_aug(input_images)
30 | is_aug = (torch.rand(batch_size, device=input_images.device)[:, None, None, None] < self.aug_proba).float()
31 | result = augmented * is_aug + permuted * (1 - is_aug)
32 | return result
33 |
34 | def _fill_masks_with_gradient(self, masks):
35 | batch_size, _, height, width = masks.shape
36 | grad = torch.linspace(0, 1, steps=width * 2, device=masks.device, dtype=masks.dtype) \
37 | .view(1, 1, 1, -1).expand(batch_size, 1, height * 2, width * 2)
38 | grad = self.grad_aug(grad)
39 | grad = CenterCrop((height, width))(grad)
40 | grad *= masks
41 |
42 | grad_for_min = grad + (1 - masks) * 10
43 | grad -= grad_for_min.view(batch_size, -1).min(-1).values[:, None, None, None]
44 | grad /= grad.view(batch_size, -1).max(-1).values[:, None, None, None] + 1e-6
45 | grad.clamp_(min=0, max=1)
46 |
47 | return grad
48 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/multidilated_conv.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import random
4 | from sd_bmab.external.lama.saicinpainting.training.modules.depthwise_sep_conv import DepthWiseSeperableConv
5 |
6 | class MultidilatedConv(nn.Module):
7 | def __init__(self, in_dim, out_dim, kernel_size, dilation_num=3, comb_mode='sum', equal_dim=True,
8 | shared_weights=False, padding=1, min_dilation=1, shuffle_in_channels=False, use_depthwise=False, **kwargs):
9 | super().__init__()
10 | convs = []
11 | self.equal_dim = equal_dim
12 | assert comb_mode in ('cat_out', 'sum', 'cat_in', 'cat_both'), comb_mode
13 | if comb_mode in ('cat_out', 'cat_both'):
14 | self.cat_out = True
15 | if equal_dim:
16 | assert out_dim % dilation_num == 0
17 | out_dims = [out_dim // dilation_num] * dilation_num
18 | self.index = sum([[i + j * (out_dims[0]) for j in range(dilation_num)] for i in range(out_dims[0])], [])
19 | else:
20 | out_dims = [out_dim // 2 ** (i + 1) for i in range(dilation_num - 1)]
21 | out_dims.append(out_dim - sum(out_dims))
22 | index = []
23 | starts = [0] + out_dims[:-1]
24 | lengths = [out_dims[i] // out_dims[-1] for i in range(dilation_num)]
25 | for i in range(out_dims[-1]):
26 | for j in range(dilation_num):
27 | index += list(range(starts[j], starts[j] + lengths[j]))
28 | starts[j] += lengths[j]
29 | self.index = index
30 | assert(len(index) == out_dim)
31 | self.out_dims = out_dims
32 | else:
33 | self.cat_out = False
34 | self.out_dims = [out_dim] * dilation_num
35 |
36 | if comb_mode in ('cat_in', 'cat_both'):
37 | if equal_dim:
38 | assert in_dim % dilation_num == 0
39 | in_dims = [in_dim // dilation_num] * dilation_num
40 | else:
41 | in_dims = [in_dim // 2 ** (i + 1) for i in range(dilation_num - 1)]
42 | in_dims.append(in_dim - sum(in_dims))
43 | self.in_dims = in_dims
44 | self.cat_in = True
45 | else:
46 | self.cat_in = False
47 | self.in_dims = [in_dim] * dilation_num
48 |
49 | conv_type = DepthWiseSeperableConv if use_depthwise else nn.Conv2d
50 | dilation = min_dilation
51 | for i in range(dilation_num):
52 | if isinstance(padding, int):
53 | cur_padding = padding * dilation
54 | else:
55 | cur_padding = padding[i]
56 | convs.append(conv_type(
57 | self.in_dims[i], self.out_dims[i], kernel_size, padding=cur_padding, dilation=dilation, **kwargs
58 | ))
59 | if i > 0 and shared_weights:
60 | convs[-1].weight = convs[0].weight
61 | convs[-1].bias = convs[0].bias
62 | dilation *= 2
63 | self.convs = nn.ModuleList(convs)
64 |
65 | self.shuffle_in_channels = shuffle_in_channels
66 | if self.shuffle_in_channels:
67 | # shuffle list as shuffling of tensors is nondeterministic
68 | in_channels_permute = list(range(in_dim))
69 | random.shuffle(in_channels_permute)
70 | # save as buffer so it is saved and loaded with checkpoint
71 | self.register_buffer('in_channels_permute', torch.tensor(in_channels_permute))
72 |
73 | def forward(self, x):
74 | if self.shuffle_in_channels:
75 | x = x[:, self.in_channels_permute]
76 |
77 | outs = []
78 | if self.cat_in:
79 | if self.equal_dim:
80 | x = x.chunk(len(self.convs), dim=1)
81 | else:
82 | new_x = []
83 | start = 0
84 | for dim in self.in_dims:
85 | new_x.append(x[:, start:start+dim])
86 | start += dim
87 | x = new_x
88 | for i, conv in enumerate(self.convs):
89 | if self.cat_in:
90 | input = x[i]
91 | else:
92 | input = x
93 | outs.append(conv(input))
94 | if self.cat_out:
95 | out = torch.cat(outs, dim=1)[:, self.index]
96 | else:
97 | out = sum(outs)
98 | return out
99 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/spatial_transform.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | from kornia.geometry.transform import rotate
5 |
6 |
7 | class LearnableSpatialTransformWrapper(nn.Module):
8 | def __init__(self, impl, pad_coef=0.5, angle_init_range=80, train_angle=True):
9 | super().__init__()
10 | self.impl = impl
11 | self.angle = torch.rand(1) * angle_init_range
12 | if train_angle:
13 | self.angle = nn.Parameter(self.angle, requires_grad=True)
14 | self.pad_coef = pad_coef
15 |
16 | def forward(self, x):
17 | if torch.is_tensor(x):
18 | return self.inverse_transform(self.impl(self.transform(x)), x)
19 | elif isinstance(x, tuple):
20 | x_trans = tuple(self.transform(elem) for elem in x)
21 | y_trans = self.impl(x_trans)
22 | return tuple(self.inverse_transform(elem, orig_x) for elem, orig_x in zip(y_trans, x))
23 | else:
24 | raise ValueError(f'Unexpected input type {type(x)}')
25 |
26 | def transform(self, x):
27 | height, width = x.shape[2:]
28 | pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)
29 | x_padded = F.pad(x, [pad_w, pad_w, pad_h, pad_h], mode='reflect')
30 | x_padded_rotated = rotate(x_padded, angle=self.angle.to(x_padded))
31 | return x_padded_rotated
32 |
33 | def inverse_transform(self, y_padded_rotated, orig_x):
34 | height, width = orig_x.shape[2:]
35 | pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)
36 |
37 | y_padded = rotate(y_padded_rotated, angle=-self.angle.to(y_padded_rotated))
38 | y_height, y_width = y_padded.shape[2:]
39 | y = y_padded[:, :, pad_h : y_height - pad_h, pad_w : y_width - pad_w]
40 | return y
41 |
42 |
43 | if __name__ == '__main__':
44 | layer = LearnableSpatialTransformWrapper(nn.Identity())
45 | x = torch.arange(2* 3 * 15 * 15).view(2, 3, 15, 15).float()
46 | y = layer(x)
47 | assert x.shape == y.shape
48 | assert torch.allclose(x[:, :, 1:, 1:][:, :, :-1, :-1], y[:, :, 1:, 1:][:, :, :-1, :-1])
49 | print('all ok')
50 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/modules/squeeze_excitation.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 |
3 |
4 | class SELayer(nn.Module):
5 | def __init__(self, channel, reduction=16):
6 | super(SELayer, self).__init__()
7 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
8 | self.fc = nn.Sequential(
9 | nn.Linear(channel, channel // reduction, bias=False),
10 | nn.ReLU(inplace=True),
11 | nn.Linear(channel // reduction, channel, bias=False),
12 | nn.Sigmoid()
13 | )
14 |
15 | def forward(self, x):
16 | b, c, _, _ = x.size()
17 | y = self.avg_pool(x).view(b, c)
18 | y = self.fc(y).view(b, c, 1, 1)
19 | res = x * y.expand_as(x)
20 | return res
21 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/trainers/__init__.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import torch
3 | from sd_bmab.external.lama.saicinpainting.training.trainers.default import DefaultInpaintingTrainingModule
4 |
5 |
6 | def get_training_model_class(kind):
7 | if kind == 'default':
8 | return DefaultInpaintingTrainingModule
9 |
10 | raise ValueError(f'Unknown trainer module {kind}')
11 |
12 |
13 | def make_training_model(config):
14 | kind = config.training_model.kind
15 | kwargs = dict(config.training_model)
16 | kwargs.pop('kind')
17 | kwargs['use_ddp'] = config.trainer.kwargs.get('accelerator', None) == 'ddp'
18 |
19 | logging.info(f'Make training model {kind}')
20 |
21 | cls = get_training_model_class(kind)
22 | return cls(config, **kwargs)
23 |
24 |
25 | def load_checkpoint(train_config, path, map_location='cuda', strict=True):
26 | model = make_training_model(train_config).generator
27 | state = torch.load(path, map_location=map_location)
28 | model.load_state_dict(state, strict=strict)
29 | return model
30 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/visualizers/__init__.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 | from sd_bmab.external.lama.saicinpainting.training.visualizers.directory import DirectoryVisualizer
4 | from sd_bmab.external.lama.saicinpainting.training.visualizers.noop import NoopVisualizer
5 |
6 |
7 | def make_visualizer(kind, **kwargs):
8 | logging.info(f'Make visualizer {kind}')
9 |
10 | if kind == 'directory':
11 | return DirectoryVisualizer(**kwargs)
12 | if kind == 'noop':
13 | return NoopVisualizer()
14 |
15 | raise ValueError(f'Unknown visualizer kind {kind}')
16 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/visualizers/base.py:
--------------------------------------------------------------------------------
1 | import abc
2 | from typing import Dict, List
3 |
4 | import numpy as np
5 | import torch
6 | from skimage import color
7 | from skimage.segmentation import mark_boundaries
8 |
9 | from . import colors
10 |
11 | COLORS, _ = colors.generate_colors(151) # 151 - max classes for semantic segmentation
12 |
13 |
14 | class BaseVisualizer:
15 | @abc.abstractmethod
16 | def __call__(self, epoch_i, batch_i, batch, suffix='', rank=None):
17 | """
18 | Take a batch, make an image from it and visualize
19 | """
20 | raise NotImplementedError()
21 |
22 |
23 | def visualize_mask_and_images(images_dict: Dict[str, np.ndarray], keys: List[str],
24 | last_without_mask=True, rescale_keys=None, mask_only_first=None,
25 | black_mask=False) -> np.ndarray:
26 | mask = images_dict['mask'] > 0.5
27 | result = []
28 | for i, k in enumerate(keys):
29 | img = images_dict[k]
30 | img = np.transpose(img, (1, 2, 0))
31 |
32 | if rescale_keys is not None and k in rescale_keys:
33 | img = img - img.min()
34 | img /= img.max() + 1e-5
35 | if len(img.shape) == 2:
36 | img = np.expand_dims(img, 2)
37 |
38 | if img.shape[2] == 1:
39 | img = np.repeat(img, 3, axis=2)
40 | elif (img.shape[2] > 3):
41 | img_classes = img.argmax(2)
42 | img = color.label2rgb(img_classes, colors=COLORS)
43 |
44 | if mask_only_first:
45 | need_mark_boundaries = i == 0
46 | else:
47 | need_mark_boundaries = i < len(keys) - 1 or not last_without_mask
48 |
49 | if need_mark_boundaries:
50 | if black_mask:
51 | img = img * (1 - mask[0][..., None])
52 | img = mark_boundaries(img,
53 | mask[0],
54 | color=(1., 0., 0.),
55 | outline_color=(1., 1., 1.),
56 | mode='thick')
57 | result.append(img)
58 | return np.concatenate(result, axis=1)
59 |
60 |
61 | def visualize_mask_and_images_batch(batch: Dict[str, torch.Tensor], keys: List[str], max_items=10,
62 | last_without_mask=True, rescale_keys=None) -> np.ndarray:
63 | batch = {k: tens.detach().cpu().numpy() for k, tens in batch.items()
64 | if k in keys or k == 'mask'}
65 |
66 | batch_size = next(iter(batch.values())).shape[0]
67 | items_to_vis = min(batch_size, max_items)
68 | result = []
69 | for i in range(items_to_vis):
70 | cur_dct = {k: tens[i] for k, tens in batch.items()}
71 | result.append(visualize_mask_and_images(cur_dct, keys, last_without_mask=last_without_mask,
72 | rescale_keys=rescale_keys))
73 | return np.concatenate(result, axis=0)
74 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/visualizers/colors.py:
--------------------------------------------------------------------------------
1 | import random
2 | import colorsys
3 |
4 | import numpy as np
5 | import matplotlib
6 | matplotlib.use('agg')
7 | import matplotlib.pyplot as plt
8 | from matplotlib.colors import LinearSegmentedColormap
9 |
10 |
11 | def generate_colors(nlabels, type='bright', first_color_black=False, last_color_black=True, verbose=False):
12 | # https://stackoverflow.com/questions/14720331/how-to-generate-random-colors-in-matplotlib
13 | """
14 | Creates a random colormap to be used together with matplotlib. Useful for segmentation tasks
15 | :param nlabels: Number of labels (size of colormap)
16 | :param type: 'bright' for strong colors, 'soft' for pastel colors
17 | :param first_color_black: Option to use first color as black, True or False
18 | :param last_color_black: Option to use last color as black, True or False
19 | :param verbose: Prints the number of labels and shows the colormap. True or False
20 | :return: colormap for matplotlib
21 | """
22 | if type not in ('bright', 'soft'):
23 | print ('Please choose "bright" or "soft" for type')
24 | return
25 |
26 | if verbose:
27 | print('Number of labels: ' + str(nlabels))
28 |
29 | # Generate color map for bright colors, based on hsv
30 | if type == 'bright':
31 | randHSVcolors = [(np.random.uniform(low=0.0, high=1),
32 | np.random.uniform(low=0.2, high=1),
33 | np.random.uniform(low=0.9, high=1)) for i in range(nlabels)]
34 |
35 | # Convert HSV list to RGB
36 | randRGBcolors = []
37 | for HSVcolor in randHSVcolors:
38 | randRGBcolors.append(colorsys.hsv_to_rgb(HSVcolor[0], HSVcolor[1], HSVcolor[2]))
39 |
40 | if first_color_black:
41 | randRGBcolors[0] = [0, 0, 0]
42 |
43 | if last_color_black:
44 | randRGBcolors[-1] = [0, 0, 0]
45 |
46 | random_colormap = LinearSegmentedColormap.from_list('new_map', randRGBcolors, N=nlabels)
47 |
48 | # Generate soft pastel colors, by limiting the RGB spectrum
49 | if type == 'soft':
50 | low = 0.6
51 | high = 0.95
52 | randRGBcolors = [(np.random.uniform(low=low, high=high),
53 | np.random.uniform(low=low, high=high),
54 | np.random.uniform(low=low, high=high)) for i in range(nlabels)]
55 |
56 | if first_color_black:
57 | randRGBcolors[0] = [0, 0, 0]
58 |
59 | if last_color_black:
60 | randRGBcolors[-1] = [0, 0, 0]
61 | random_colormap = LinearSegmentedColormap.from_list('new_map', randRGBcolors, N=nlabels)
62 |
63 | # Display colorbar
64 | if verbose:
65 | from matplotlib import colors, colorbar
66 | from matplotlib import pyplot as plt
67 | fig, ax = plt.subplots(1, 1, figsize=(15, 0.5))
68 |
69 | bounds = np.linspace(0, nlabels, nlabels + 1)
70 | norm = colors.BoundaryNorm(bounds, nlabels)
71 |
72 | cb = colorbar.ColorbarBase(ax, cmap=random_colormap, norm=norm, spacing='proportional', ticks=None,
73 | boundaries=bounds, format='%1i', orientation=u'horizontal')
74 |
75 | return randRGBcolors, random_colormap
76 |
77 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/visualizers/directory.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | import cv2
4 | import numpy as np
5 |
6 | from sd_bmab.external.lama.saicinpainting.training.visualizers.base import BaseVisualizer, visualize_mask_and_images_batch
7 | from sd_bmab.external.lama.saicinpainting.utils import check_and_warn_input_range
8 |
9 |
10 | class DirectoryVisualizer(BaseVisualizer):
11 | DEFAULT_KEY_ORDER = 'image predicted_image inpainted'.split(' ')
12 |
13 | def __init__(self, outdir, key_order=DEFAULT_KEY_ORDER, max_items_in_batch=10,
14 | last_without_mask=True, rescale_keys=None):
15 | self.outdir = outdir
16 | os.makedirs(self.outdir, exist_ok=True)
17 | self.key_order = key_order
18 | self.max_items_in_batch = max_items_in_batch
19 | self.last_without_mask = last_without_mask
20 | self.rescale_keys = rescale_keys
21 |
22 | def __call__(self, epoch_i, batch_i, batch, suffix='', rank=None):
23 | check_and_warn_input_range(batch['image'], 0, 1, 'DirectoryVisualizer target image')
24 | vis_img = visualize_mask_and_images_batch(batch, self.key_order, max_items=self.max_items_in_batch,
25 | last_without_mask=self.last_without_mask,
26 | rescale_keys=self.rescale_keys)
27 |
28 | vis_img = np.clip(vis_img * 255, 0, 255).astype('uint8')
29 |
30 | curoutdir = os.path.join(self.outdir, f'epoch{epoch_i:04d}{suffix}')
31 | os.makedirs(curoutdir, exist_ok=True)
32 | rank_suffix = f'_r{rank}' if rank is not None else ''
33 | out_fname = os.path.join(curoutdir, f'batch{batch_i:07d}{rank_suffix}.jpg')
34 |
35 | vis_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
36 | cv2.imwrite(out_fname, vis_img)
37 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/training/visualizers/noop.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.external.lama.saicinpainting.training.visualizers.base import BaseVisualizer
2 |
3 |
4 | class NoopVisualizer(BaseVisualizer):
5 | def __init__(self, *args, **kwargs):
6 | pass
7 |
8 | def __call__(self, epoch_i, batch_i, batch, suffix='', rank=None):
9 | pass
10 |
--------------------------------------------------------------------------------
/sd_bmab/external/lama/saicinpainting/utils.py:
--------------------------------------------------------------------------------
1 | import bisect
2 | import functools
3 | import logging
4 | import numbers
5 | import os
6 | import signal
7 | import sys
8 | import traceback
9 | import warnings
10 |
11 | import torch
12 | from pytorch_lightning import seed_everything
13 |
14 | LOGGER = logging.getLogger(__name__)
15 |
16 |
17 | def check_and_warn_input_range(tensor, min_value, max_value, name):
18 | actual_min = tensor.min()
19 | actual_max = tensor.max()
20 | if actual_min < min_value or actual_max > max_value:
21 | warnings.warn(f"{name} must be in {min_value}..{max_value} range, but it ranges {actual_min}..{actual_max}")
22 |
23 |
24 | def sum_dict_with_prefix(target, cur_dict, prefix, default=0):
25 | for k, v in cur_dict.items():
26 | target_key = prefix + k
27 | target[target_key] = target.get(target_key, default) + v
28 |
29 |
30 | def average_dicts(dict_list):
31 | result = {}
32 | norm = 1e-3
33 | for dct in dict_list:
34 | sum_dict_with_prefix(result, dct, '')
35 | norm += 1
36 | for k in list(result):
37 | result[k] /= norm
38 | return result
39 |
40 |
41 | def add_prefix_to_keys(dct, prefix):
42 | return {prefix + k: v for k, v in dct.items()}
43 |
44 |
45 | def set_requires_grad(module, value):
46 | for param in module.parameters():
47 | param.requires_grad = value
48 |
49 |
50 | def flatten_dict(dct):
51 | result = {}
52 | for k, v in dct.items():
53 | if isinstance(k, tuple):
54 | k = '_'.join(k)
55 | if isinstance(v, dict):
56 | for sub_k, sub_v in flatten_dict(v).items():
57 | result[f'{k}_{sub_k}'] = sub_v
58 | else:
59 | result[k] = v
60 | return result
61 |
62 |
63 | class LinearRamp:
64 | def __init__(self, start_value=0, end_value=1, start_iter=-1, end_iter=0):
65 | self.start_value = start_value
66 | self.end_value = end_value
67 | self.start_iter = start_iter
68 | self.end_iter = end_iter
69 |
70 | def __call__(self, i):
71 | if i < self.start_iter:
72 | return self.start_value
73 | if i >= self.end_iter:
74 | return self.end_value
75 | part = (i - self.start_iter) / (self.end_iter - self.start_iter)
76 | return self.start_value * (1 - part) + self.end_value * part
77 |
78 |
79 | class LadderRamp:
80 | def __init__(self, start_iters, values):
81 | self.start_iters = start_iters
82 | self.values = values
83 | assert len(values) == len(start_iters) + 1, (len(values), len(start_iters))
84 |
85 | def __call__(self, i):
86 | segment_i = bisect.bisect_right(self.start_iters, i)
87 | return self.values[segment_i]
88 |
89 |
90 | def get_ramp(kind='ladder', **kwargs):
91 | if kind == 'linear':
92 | return LinearRamp(**kwargs)
93 | if kind == 'ladder':
94 | return LadderRamp(**kwargs)
95 | raise ValueError(f'Unexpected ramp kind: {kind}')
96 |
97 |
98 | def print_traceback_handler(sig, frame):
99 | LOGGER.warning(f'Received signal {sig}')
100 | bt = ''.join(traceback.format_stack())
101 | LOGGER.warning(f'Requested stack trace:\n{bt}')
102 |
103 |
104 | def register_debug_signal_handlers(sig=None, handler=print_traceback_handler):
105 | LOGGER.warning(f'Setting signal {sig} handler {handler}')
106 | signal.signal(sig, handler)
107 |
108 |
109 | def handle_deterministic_config(config):
110 | seed = dict(config).get('seed', None)
111 | if seed is None:
112 | return False
113 |
114 | seed_everything(seed)
115 | return True
116 |
117 |
118 | def get_shape(t):
119 | if torch.is_tensor(t):
120 | return tuple(t.shape)
121 | elif isinstance(t, dict):
122 | return {n: get_shape(q) for n, q in t.items()}
123 | elif isinstance(t, (list, tuple)):
124 | return [get_shape(q) for q in t]
125 | elif isinstance(t, numbers.Number):
126 | return type(t)
127 | else:
128 | raise ValueError('unexpected type {}'.format(type(t)))
129 |
130 |
131 | def get_has_ddp_rank():
132 | master_port = os.environ.get('MASTER_PORT', None)
133 | node_rank = os.environ.get('NODE_RANK', None)
134 | local_rank = os.environ.get('LOCAL_RANK', None)
135 | world_size = os.environ.get('WORLD_SIZE', None)
136 | has_rank = master_port is not None or node_rank is not None or local_rank is not None or world_size is not None
137 | return has_rank
138 |
139 |
140 | def handle_ddp_subprocess():
141 | def main_decorator(main_func):
142 | @functools.wraps(main_func)
143 | def new_main(*args, **kwargs):
144 | # Trainer sets MASTER_PORT, NODE_RANK, LOCAL_RANK, WORLD_SIZE
145 | parent_cwd = os.environ.get('TRAINING_PARENT_WORK_DIR', None)
146 | has_parent = parent_cwd is not None
147 | has_rank = get_has_ddp_rank()
148 | assert has_parent == has_rank, f'Inconsistent state: has_parent={has_parent}, has_rank={has_rank}'
149 |
150 | if has_parent:
151 | # we are in the worker
152 | sys.argv.extend([
153 | f'hydra.run.dir={parent_cwd}',
154 | # 'hydra/hydra_logging=disabled',
155 | # 'hydra/job_logging=disabled'
156 | ])
157 | # do nothing if this is a top-level process
158 | # TRAINING_PARENT_WORK_DIR is set in handle_ddp_parent_process after hydra initialization
159 |
160 | main_func(*args, **kwargs)
161 | return new_main
162 | return main_decorator
163 |
164 |
165 | def handle_ddp_parent_process():
166 | parent_cwd = os.environ.get('TRAINING_PARENT_WORK_DIR', None)
167 | has_parent = parent_cwd is not None
168 | has_rank = get_has_ddp_rank()
169 | assert has_parent == has_rank, f'Inconsistent state: has_parent={has_parent}, has_rank={has_rank}'
170 |
171 | if parent_cwd is None:
172 | os.environ['TRAINING_PARENT_WORK_DIR'] = os.getcwd()
173 |
174 | return has_parent
175 |
--------------------------------------------------------------------------------
/sd_bmab/masking/__init__.py:
--------------------------------------------------------------------------------
1 | from modules import shared
2 |
3 | from sd_bmab.base import MaskBase
4 | from sd_bmab.masking.sam import SamPredictVitB
5 | from sd_bmab.masking.sam_hq import SamHqPredictVitB
6 | from sd_bmab.util import debug_print
7 |
8 |
9 | masks = [SamPredictVitB(), SamHqPredictVitB()]
10 | dmasks = {s.name: s for s in masks}
11 |
12 |
13 | def get_mask_generator(name='None') -> MaskBase:
14 | model = dmasks.get(name, dmasks[shared.opts.bmab_mask_model])
15 | debug_print(f'Use mask model {model.name}')
16 | return model
17 |
18 |
19 | def list_mask_names():
20 | return [s.name for s in masks]
21 |
22 |
23 | def release():
24 | SamPredictVitB.release()
25 | SamHqPredictVitB.release()
26 |
--------------------------------------------------------------------------------
/sd_bmab/masking/sam.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base import MaskBase
2 | import cv2
3 | import numpy as np
4 |
5 | import torch
6 |
7 | from PIL import Image
8 | from modules.safe import unsafe_torch_load, load
9 | from modules.devices import device, torch_gc
10 |
11 | from segment_anything import SamPredictor
12 | from segment_anything import sam_model_registry
13 |
14 | from sd_bmab import util
15 |
16 |
17 | sam_model = None
18 |
19 |
20 | class SamPredict(MaskBase):
21 |
22 | def __init__(self) -> None:
23 | super().__init__()
24 |
25 | @property
26 | def name(self):
27 | return f'sam_{self.type}'
28 |
29 | @property
30 | def type(self):
31 | pass
32 |
33 | @property
34 | def file(self):
35 | pass
36 |
37 | @classmethod
38 | def init(cls, model_type, filename, *args, **kwargs):
39 | checkpoint_file = util.lazy_loader(filename)
40 | global sam_model
41 | if sam_model is None:
42 | torch.load = unsafe_torch_load
43 | sam_model = sam_model_registry[model_type](checkpoint=checkpoint_file)
44 | sam_model.to(device=device)
45 | torch.load = load
46 | return sam_model
47 |
48 | def load(self):
49 | return SamPredict.init(self.type, self.file)
50 |
51 | def predict(self, image, boxes):
52 | sam = self.load()
53 | mask_predictor = SamPredictor(sam)
54 |
55 | numpy_image = np.array(image)
56 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
57 | mask_predictor.set_image(opencv_image)
58 |
59 | result = Image.new('L', image.size, 0)
60 | for box in boxes:
61 | x1, y1, x2, y2 = box
62 |
63 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
64 | masks, scores, logits = mask_predictor.predict(
65 | box=box,
66 | multimask_output=False
67 | )
68 |
69 | mask = Image.fromarray(masks[0])
70 | result.paste(mask, mask=mask)
71 | return result
72 |
73 | def predict_multiple(self, image, points, labels, boxes=None):
74 | sam = self.load()
75 | mask_predictor = SamPredictor(sam)
76 |
77 | numpy_image = np.array(image)
78 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
79 | mask_predictor.set_image(opencv_image)
80 |
81 | result = Image.new('L', image.size, 0)
82 | kwargs = dict(multimask_output=True)
83 | if len(points) > 0:
84 | kwargs['point_coords'] = points
85 | kwargs['point_labels'] = labels
86 | if boxes is not None and boxes:
87 | for box in boxes:
88 | x1, y1, x2, y2 = box
89 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
90 | masks, scores, logits = mask_predictor.predict(
91 | box=box,
92 | **kwargs
93 | )
94 | for mask in masks:
95 | mask = Image.fromarray(mask)
96 | result.paste(mask, mask=mask)
97 | else:
98 | masks, scores, logits = mask_predictor.predict(
99 | **kwargs
100 | )
101 | for mask in masks:
102 | mask = Image.fromarray(mask)
103 | result.paste(mask, mask=mask)
104 | return result
105 |
106 | @classmethod
107 | def release(cls):
108 | global sam_model
109 | if sam_model is not None:
110 | sam_model = None
111 | torch_gc()
112 |
113 |
114 | class SamPredictVitB(SamPredict):
115 |
116 | @property
117 | def type(self):
118 | return 'vit_b'
119 |
120 | @property
121 | def file(self):
122 | return 'sam_vit_b_01ec64.pth'
123 |
--------------------------------------------------------------------------------
/sd_bmab/masking/sam_hq.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base import MaskBase
2 | import cv2
3 | import numpy as np
4 |
5 | import torch
6 |
7 | from PIL import Image
8 | from modules.safe import unsafe_torch_load, load
9 | from modules.devices import device, torch_gc
10 |
11 | from segment_anything_hq import SamPredictor
12 | from segment_anything_hq import sam_model_registry
13 |
14 | from sd_bmab import util
15 |
16 |
17 | sam_hq_model = None
18 |
19 |
20 | class SamHqPredict(MaskBase):
21 |
22 | def __init__(self) -> None:
23 | super().__init__()
24 |
25 | @property
26 | def name(self):
27 | return f'sam_hq_{self.type}'
28 |
29 | @property
30 | def type(self):
31 | pass
32 |
33 | @property
34 | def file(self):
35 | pass
36 |
37 | @classmethod
38 | def init(cls, model_type, filename, *args, **kwargs):
39 | checkpoint_file = util.lazy_loader(filename)
40 | global sam_hq_model
41 | if sam_hq_model is None:
42 | torch.load = unsafe_torch_load
43 | sam_hq_model = sam_model_registry[model_type](checkpoint=checkpoint_file)
44 | sam_hq_model.to(device=device)
45 | torch.load = load
46 | return sam_hq_model
47 |
48 | def load(self):
49 | return SamHqPredict.init(self.type, self.file)
50 |
51 | def predict(self, image, boxes):
52 | sam = self.load()
53 | mask_predictor = SamPredictor(sam)
54 |
55 | numpy_image = np.array(image)
56 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
57 | mask_predictor.set_image(opencv_image)
58 |
59 | result = Image.new('L', image.size, 0)
60 | for box in boxes:
61 | x1, y1, x2, y2 = box
62 |
63 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
64 | masks, scores, logits = mask_predictor.predict(
65 | box=box,
66 | multimask_output=False
67 | )
68 |
69 | mask = Image.fromarray(masks[0])
70 | result.paste(mask, mask=mask)
71 | return result
72 |
73 | def predict_multiple(self, image, points, labels, boxes=None):
74 | sam = self.load()
75 | mask_predictor = SamPredictor(sam)
76 |
77 | numpy_image = np.array(image)
78 | opencv_image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
79 | mask_predictor.set_image(opencv_image)
80 |
81 | result = Image.new('L', image.size, 0)
82 | kwargs = dict(multimask_output=True)
83 | if len(points) > 0:
84 | kwargs['point_coords'] = points
85 | kwargs['point_labels'] = labels
86 | if boxes is not None:
87 | for box in boxes:
88 | x1, y1, x2, y2 = box
89 | box = np.array([int(x1), int(y1), int(x2), int(y2)])
90 | masks, scores, logits = mask_predictor.predict(
91 | box=box,
92 | **kwargs
93 | )
94 | for mask in masks:
95 | mask = Image.fromarray(mask)
96 | result.paste(mask, mask=mask)
97 | else:
98 | masks, scores, logits = mask_predictor.predict(
99 | **kwargs
100 | )
101 | for mask in masks:
102 | mask = Image.fromarray(mask)
103 | result.paste(mask, mask=mask)
104 | return result
105 |
106 | @classmethod
107 | def release(cls):
108 | global sam_hq_model
109 | if sam_hq_model is not None:
110 | sam_hq_model = None
111 | torch_gc()
112 |
113 |
114 | class SamHqPredictVitB(SamHqPredict):
115 |
116 | @property
117 | def type(self):
118 | return 'vit_b'
119 |
120 | @property
121 | def file(self):
122 | return 'sam_hq_vit_b.pth'
123 |
--------------------------------------------------------------------------------
/sd_bmab/pipeline/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/pipeline/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/pipeline/internal/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.pipeline.internal.internalpipe import process_img2img, process_intermediate_before_upscale, process_intermediate_after_upscale, is_controlnet_required
2 | from sd_bmab.pipeline.internal.intermediate import Preprocess
3 |
--------------------------------------------------------------------------------
/sd_bmab/pipeline/internal/intermediate.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from sd_bmab.base.context import Context
4 | from sd_bmab.base.processorbase import ProcessorBase
5 | from sd_bmab.pipeline.internal import process_intermediate_before_upscale, process_intermediate_after_upscale
6 |
7 |
8 | class Preprocess(ProcessorBase):
9 | def __init__(self) -> None:
10 | super().__init__()
11 |
12 | def preprocess(self, context: Context, image: Image):
13 | return context.is_txtimg() and not context.is_hires_fix()
14 |
15 | def process(self, context: Context, image: Image):
16 | image = process_intermediate_before_upscale(context, image)
17 | return process_intermediate_after_upscale(context, image)
18 |
19 | def postprocess(self, context: Context, image: Image):
20 | pass
21 |
--------------------------------------------------------------------------------
/sd_bmab/pipeline/internal/internalpipe.py:
--------------------------------------------------------------------------------
1 | import traceback
2 | from functools import partial
3 | from modules import images
4 | from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img
5 | from modules import shared
6 |
7 | from sd_bmab.util import debug_print
8 | from sd_bmab.base import Context
9 | from sd_bmab.processors.detailer import FaceDetailer, FaceDetailerBeforeUpsacle
10 |
11 | from sd_bmab.processors.basic import EdgeEnhancement, NoiseAlpha, Img2imgMasking, ICLightBeforeUpsacle
12 | from sd_bmab.processors.preprocess import ResizeIntermidiateBeforeUpscale
13 | from sd_bmab.processors.preprocess import ResamplePreprocessorBeforeUpscale, ResizeIntermidiateAfterUpsacle
14 | from sd_bmab.processors.preprocess import PretrainingDetailerBeforeUpscale
15 |
16 |
17 | def is_controlnet_required(context):
18 | pipeline_modules = [
19 | ResamplePreprocessorBeforeUpscale(),
20 | ResizeIntermidiateBeforeUpscale()
21 | ]
22 | for mod in pipeline_modules:
23 | if mod.use_controlnet(context):
24 | return True
25 | return False
26 |
27 |
28 | def process_intermediate_before_upscale(context, image):
29 | pipeline_before_upscale = [
30 | ResamplePreprocessorBeforeUpscale(),
31 | PretrainingDetailerBeforeUpscale(),
32 | FaceDetailerBeforeUpsacle(),
33 | ICLightBeforeUpsacle(),
34 | ResizeIntermidiateBeforeUpscale(),
35 | ]
36 |
37 | processed = image.copy()
38 | for proc in pipeline_before_upscale:
39 | try:
40 | result = proc.preprocess(context, processed)
41 | if result is None or not result:
42 | continue
43 | if shared.state.interrupted or shared.state.skipped:
44 | break
45 | ret = proc.process(context, processed)
46 | if shared.state.interrupted or shared.state.skipped:
47 | break
48 | proc.postprocess(context, processed)
49 | processed = ret
50 | except Exception:
51 | traceback.print_exc()
52 | return processed
53 |
54 |
55 | def process_intermediate_after_upscale(context, image):
56 | pipeline_before_upscale = [
57 | EdgeEnhancement(),
58 | ResizeIntermidiateAfterUpsacle(),
59 | Img2imgMasking(),
60 | NoiseAlpha(),
61 | ]
62 |
63 | processed = image.copy()
64 | for proc in pipeline_before_upscale:
65 | result = proc.preprocess(context, processed)
66 | if result is None or not result:
67 | continue
68 | if shared.state.interrupted or shared.state.skipped:
69 | break
70 | ret = proc.process(context, processed)
71 | proc.postprocess(context, processed)
72 | processed = ret
73 | return processed
74 |
75 |
76 | def process_img2img(ctx):
77 | if not ctx.is_img2img():
78 | return
79 |
80 | image = ctx.sdprocessing.init_images[0]
81 | debug_print('process img2img ', image.size)
82 | image = process_intermediate_before_upscale(ctx, image)
83 | image = process_intermediate_after_upscale(ctx, image)
84 | debug_print('process img2img ', image.size)
85 | ctx.sdprocessing.init_images[0] = image
86 |
87 |
88 | '''
89 | def process_hiresfix(ctx):
90 | if not isinstance(ctx.sdprocessing, StableDiffusionProcessingTxt2Img):
91 | return
92 |
93 | if hasattr(ctx.sdprocessing, '__sample'):
94 | return
95 |
96 | all_processors = [
97 | FaceDetailer(),
98 | EdgeEnhancement(),
99 | ResizeIntermidiateBeforeUpscale(),
100 | ResizeIntermidiateAfterUpsacle(),
101 | NoiseAlpha()
102 | ]
103 |
104 | if True not in [proc.preprocess(ctx, None) for proc in all_processors]:
105 | return
106 |
107 | ctx.sdprocessing.__sample = ctx.sdprocessing.sample
108 |
109 | def resize(ctx: Context, resize_mode, img, width, height, upscaler_name=None):
110 | images.resize_image = ctx.sdprocessing.resize_hook
111 | pidx = ctx.sdprocessing.iteration * ctx.sdprocessing.batch_size
112 | ctx.index += 1
113 | ctx.args['current_prompt'] = ctx.sdprocessing.all_prompts[pidx]
114 | img = process_intermediate_step1(ctx, img)
115 | im = ctx.sdprocessing.resize_hook(resize_mode, img, width, height, upscaler_name)
116 | im = process_intermediate_step2(ctx, im)
117 | images.resize_image = partial(resize, ctx)
118 | return im
119 |
120 | def _sample(ctx: Context, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
121 | ctx.sdprocessing.resize_hook = images.resize_image
122 | images.resize_image = partial(resize, ctx)
123 | try:
124 | ret = ctx.sdprocessing.__sample(conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts)
125 | except Exception as e:
126 | raise e
127 | finally:
128 | images.resize_image = ctx.sdprocessing.resize_hook
129 | return ret
130 |
131 | ctx.sdprocessing.sample = partial(_sample, ctx)
132 | '''
133 |
--------------------------------------------------------------------------------
/sd_bmab/pipeline/post/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.pipeline.post.mainpipe import process, process_controlnet, release, is_controlnet_required
2 |
--------------------------------------------------------------------------------
/sd_bmab/pipeline/post/mainpipe.py:
--------------------------------------------------------------------------------
1 | import traceback
2 |
3 | from modules import shared
4 |
5 | from sd_bmab.processors.postprocess import AfterProcessUpscaler, BeforeProcessUpscaler
6 | from sd_bmab.processors.postprocess import InpaintResize, InpaintLamaResize, FinalFilter
7 | from sd_bmab.processors.detailer import FaceDetailer, PersonDetailer, HandDetailer, PreprocessFaceDetailer
8 | from sd_bmab.processors.utils import BeforeProcessFileSaver, AfterProcessFileSaver
9 | from sd_bmab.processors.utils import ApplyModel, RollbackModel, CheckPointChanger, CheckPointRestore
10 | from sd_bmab.processors.basic import FinalProcessorBasic, ICLightAfterUpsacle
11 | from sd_bmab.processors.controlnet import LineartNoise, Openpose, IpAdapter
12 | from sd_bmab.processors.preprocess import RefinerPreprocessor, PretrainingDetailer
13 | from sd_bmab.processors.preprocess import ResamplePreprocessor, PreprocessFilter
14 | from sd_bmab.processors.postprocess import Watermark
15 | from sd_bmab.pipeline.internal import Preprocess
16 | from sd_bmab.util import debug_print
17 |
18 |
19 | def is_controlnet_required(context):
20 | pipeline_modules = [
21 | ResamplePreprocessor(),
22 | InpaintLamaResize(),
23 | ]
24 | for mod in pipeline_modules:
25 | if mod.use_controlnet(context):
26 | return True
27 | return False
28 |
29 |
30 | def process(context, image):
31 | pipeline_modules = [
32 | BeforeProcessFileSaver(),
33 | PreprocessFaceDetailer(),
34 | CheckPointChanger(),
35 | PreprocessFilter(),
36 | ResamplePreprocessor(),
37 | PretrainingDetailer(),
38 | Preprocess(),
39 | RefinerPreprocessor(),
40 | ICLightAfterUpsacle(),
41 | InpaintResize(),
42 | InpaintLamaResize(),
43 | BeforeProcessUpscaler(),
44 | ApplyModel(),
45 | PersonDetailer(),
46 | FaceDetailer(),
47 | HandDetailer(),
48 | RollbackModel(),
49 | CheckPointRestore(),
50 | AfterProcessUpscaler(),
51 | FinalProcessorBasic(),
52 | FinalFilter(),
53 | Watermark(),
54 | AfterProcessFileSaver()
55 | ]
56 |
57 | processed = image.copy()
58 |
59 | try:
60 | for proc in pipeline_modules:
61 | try:
62 | if shared.state.interrupted or shared.state.skipped:
63 | return
64 |
65 | result = proc.preprocess(context, processed)
66 | if result is None or not result:
67 | continue
68 | if shared.state.interrupted or shared.state.skipped:
69 | break
70 | ret = proc.process(context, processed)
71 | if shared.state.interrupted or shared.state.skipped:
72 | break
73 | proc.postprocess(context, processed)
74 | processed = ret
75 | except:
76 | traceback.print_exc()
77 | finally:
78 | debug_print('Restore Checkpoint at final')
79 | RollbackModel().process(context, processed)
80 | CheckPointRestore().process(context, processed)
81 | for proc in pipeline_modules:
82 | try:
83 | proc.finalprocess(context, processed)
84 | except:
85 | traceback.print_exc()
86 | return processed
87 |
88 | '''
89 | def process_intermediate(context, image):
90 | all_processors = [
91 | FaceDetailer(),
92 | EdgeEnhancement(),
93 | ResizeIntermidiate(),
94 | NoiseAlpha()
95 | ]
96 |
97 | processed = image.copy()
98 |
99 | for proc in all_processors:
100 | result = proc.preprocess(context, processed)
101 | if result is None or not result:
102 | continue
103 | ret = proc.process(context, processed)
104 | proc.postprocess(context, processed)
105 | processed = ret
106 |
107 | return processed
108 | '''
109 |
110 | def process_controlnet(context):
111 | all_processors = [
112 | LineartNoise(),
113 | Openpose(),
114 | IpAdapter()
115 | ]
116 |
117 | for proc in all_processors:
118 | result = proc.preprocess(context, None)
119 | if result is None or not result:
120 | continue
121 | proc.process(context, None)
122 | proc.postprocess(context, None)
123 |
124 |
125 | def release():
126 | pass
127 |
--------------------------------------------------------------------------------
/sd_bmab/processors/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/portu-sim/sd-webui-bmab/7a81441a72cef95cd9fdb732ee933bb04799dab3/sd_bmab/processors/__init__.py
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.basic.edge import EdgeEnhancement
2 | from sd_bmab.processors.basic.final import FinalProcessorBasic
3 | from sd_bmab.processors.basic.intermediate import NoiseAlpha
4 | from sd_bmab.processors.basic.blend import BlendImage
5 | from sd_bmab.processors.basic.img2imgmasking import Img2imgMasking
6 | from sd_bmab.processors.basic.iclight import ICLight, ICLightBeforeUpsacle, ICLightAfterUpsacle
7 |
8 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/blend.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from sd_bmab.base.context import Context
4 | from sd_bmab.base.processorbase import ProcessorBase
5 | from sd_bmab.util import debug_print
6 |
7 |
8 | class BlendImage(ProcessorBase):
9 | def __init__(self) -> None:
10 | super().__init__()
11 | self.enabled = False
12 | self.input_image = None
13 | self.alpha = 0
14 |
15 | def preprocess(self, context: Context, image: Image):
16 | self.enabled = context.args['blend_enabled']
17 | self.input_image = context.args['input_image']
18 | self.alpha = context.args['blend_alpha']
19 | #return self.enabled and self.input_image is not None and 0 <= self.alpha <= 1
20 | return False
21 |
22 | def process(self, context: Context, image: Image):
23 | context.add_generation_param('BMAB blend alpha', self.alpha)
24 | #blend = Image.fromarray(self.input_image, mode='RGB')
25 | blend = self.input_image
26 | img = Image.new(mode='RGB', size=image.size)
27 | img.paste(image, (0, 0))
28 | img.paste(blend)
29 | image = Image.blend(image, img, alpha=self.alpha)
30 | return image
31 |
32 | def postprocess(self, context: Context, image: Image):
33 | pass
34 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/edge.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 |
4 | from PIL import Image
5 | from PIL import ImageOps
6 | from sd_bmab.base.context import Context
7 | from sd_bmab.base.processorbase import ProcessorBase
8 |
9 |
10 | class EdgeEnhancement(ProcessorBase):
11 | def __init__(self) -> None:
12 | super().__init__()
13 |
14 | self.edge_low_threadhold = 50
15 | self.edge_high_threadhold = 200
16 | self.edge_strength = 0.5
17 |
18 | def preprocess(self, context: Context, image: Image):
19 | if context.args['edge_flavor_enabled']:
20 | self.edge_low_threadhold = context.args['edge_low_threadhold']
21 | self.edge_high_threadhold = context.args['edge_high_threadhold']
22 | self.edge_strength = context.args['edge_strength']
23 | return context.args['edge_flavor_enabled']
24 |
25 | def process(self, context: Context, image: Image):
26 | context.add_generation_param('BMAB edge flavor low threadhold', self.edge_low_threadhold)
27 | context.add_generation_param('BMAB edge flavor high threadhold', self.edge_high_threadhold)
28 | context.add_generation_param('BMAB edge flavor strength', self.edge_strength)
29 |
30 | numpy_image = np.array(image)
31 | base = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
32 | arcanny = cv2.Canny(base, self.edge_low_threadhold, self.edge_high_threadhold)
33 | canny = Image.fromarray(arcanny)
34 | canny = ImageOps.invert(canny)
35 |
36 | img = image.convert('RGB')
37 | newdata = [(0, 0, 0) if mdata == 0 else ndata for mdata, ndata in zip(canny.getdata(), img.getdata())]
38 | newbase = Image.new('RGB', img.size)
39 | newbase.putdata(newdata)
40 | return Image.blend(img, newbase, alpha=self.edge_strength).convert("RGB")
41 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/final.py:
--------------------------------------------------------------------------------
1 | import math
2 |
3 | from PIL import Image
4 | from PIL import ImageEnhance
5 |
6 | from sd_bmab import util
7 | from sd_bmab.base.context import Context
8 | from sd_bmab.base.processorbase import ProcessorBase
9 |
10 |
11 | def calc_color_temperature(temp):
12 | white = (255.0, 254.11008387561782, 250.0419083427406)
13 |
14 | temperature = temp / 100
15 |
16 | if temperature <= 66:
17 | red = 255.0
18 | else:
19 | red = float(temperature - 60)
20 | red = 329.698727446 * math.pow(red, -0.1332047592)
21 | if red < 0:
22 | red = 0
23 | if red > 255:
24 | red = 255
25 |
26 | if temperature <= 66:
27 | green = temperature
28 | green = 99.4708025861 * math.log(green) - 161.1195681661
29 | else:
30 | green = float(temperature - 60)
31 | green = 288.1221695283 * math.pow(green, -0.0755148492)
32 | if green < 0:
33 | green = 0
34 | if green > 255:
35 | green = 255
36 |
37 | if temperature >= 66:
38 | blue = 255.0
39 | else:
40 | if temperature <= 19:
41 | blue = 0.0
42 | else:
43 | blue = float(temperature - 10)
44 | blue = 138.5177312231 * math.log(blue) - 305.0447927307
45 | if blue < 0:
46 | blue = 0
47 | if blue > 255:
48 | blue = 255
49 |
50 | return red / white[0], green / white[1], blue / white[2]
51 |
52 |
53 | class FinalProcessorBasic(ProcessorBase):
54 | def __init__(self) -> None:
55 | super().__init__()
56 | self.noise_alpha_final = 0
57 | self.contrast = 1
58 | self.brightness = 1
59 | self.sharpeness = 1
60 | self.color_saturation = 1
61 | self.color_temperature = 0
62 |
63 | def preprocess(self, context: Context, image: Image):
64 | self.contrast = context.args['contrast']
65 | self.brightness = context.args['brightness']
66 | self.sharpeness = context.args['sharpeness']
67 | self.color_saturation = context.args['color_saturation']
68 | self.color_temperature = context.args['color_temperature']
69 | self.noise_alpha_final = context.args['noise_alpha_final']
70 | return True
71 |
72 | def process(self, context: Context, image: Image):
73 |
74 | if self.noise_alpha_final != 0:
75 | context.add_generation_param('BMAB noise alpha final', self.noise_alpha_final)
76 | img_noise = util.generate_noise(context.sdprocessing.seed, image.size[0], image.size[1])
77 | image = Image.blend(image, img_noise, alpha=self.noise_alpha_final)
78 |
79 | if self.contrast != 1:
80 | context.add_generation_param('BMAB contrast', self.contrast)
81 | enhancer = ImageEnhance.Contrast(image)
82 | image = enhancer.enhance(self.contrast)
83 |
84 | if self.brightness != 1:
85 | context.add_generation_param('BMAB brightness', self.brightness)
86 | enhancer = ImageEnhance.Brightness(image)
87 | image = enhancer.enhance(self.brightness)
88 |
89 | if self.sharpeness != 1:
90 | context.add_generation_param('BMAB sharpeness', self.sharpeness)
91 | enhancer = ImageEnhance.Sharpness(image)
92 | image = enhancer.enhance(self.sharpeness)
93 |
94 | if self.color_saturation != 1:
95 | context.add_generation_param('BMAB color', self.color_saturation)
96 | enhancer = ImageEnhance.Color(image)
97 | image = enhancer.enhance(self.color_saturation)
98 |
99 | if self.color_temperature != 0:
100 | context.add_generation_param('BMAB color temperature', self.color_temperature)
101 | temp = calc_color_temperature(6500 + self.color_temperature)
102 | az = []
103 | for d in image.getdata():
104 | az.append((int(d[0] * temp[0]), int(d[1] * temp[1]), int(d[2] * temp[2])))
105 | image = Image.new('RGB', image.size)
106 | image.putdata(az)
107 |
108 | return image
109 |
110 | def postprocess(self, context: Context, image: Image):
111 | pass
112 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/iclight.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from sd_bmab.base import cache
4 | from sd_bmab.base import ProcessorBase
5 | from sd_bmab.base.context import Context
6 | from sd_bmab.external import load_external_module
7 |
8 |
9 | class ICLight(ProcessorBase):
10 |
11 | def __init__(self) -> None:
12 | super().__init__()
13 | self.iclight_opt = {}
14 | self.enabled = False
15 | self.enable_before_upscale = False
16 | self.style = 'normal'
17 | self.prompt = ''
18 | self.preference = 'None'
19 | self.use_background_image = False
20 | self.blending = 0.5
21 |
22 | def preprocess(self, context: Context, image: Image):
23 | self.iclight_opt = context.args.get('module_config', {}).get('iclight', {})
24 | self.enabled = self.iclight_opt.get('enabled', self.enabled)
25 | self.enable_before_upscale = self.iclight_opt.get('enable_before_upscale', self.enable_before_upscale)
26 | self.style = self.iclight_opt.get('style', self.style)
27 | self.prompt = self.iclight_opt.get('prompt', self.prompt)
28 | self.preference = self.iclight_opt.get('preference', self.preference)
29 | self.use_background_image = self.iclight_opt.get('use_background_image', self.use_background_image)
30 | self.blending = self.iclight_opt.get('blending', self.blending)
31 | return self.enabled
32 |
33 | def process(self, context: Context, image: Image):
34 | mod = load_external_module('iclight', 'bmabiclight')
35 | bg_image = self.get_background_image() if self.use_background_image else None
36 | return mod.bmab_relight(context, self.style, image, bg_image, self.prompt, self.blending, self.preference)
37 |
38 | def postprocess(self, context: Context, image: Image):
39 | pass
40 |
41 | @staticmethod
42 | def get_background_image():
43 | img = cache.get_image_from_cache('iclight_background.png')
44 | if img is not None:
45 | return img
46 | return Image.new('RGB', (512, 768), (0, 0, 0))
47 |
48 | @staticmethod
49 | def put_backgound_image(img):
50 | cache.put_image_to_cache('iclight_background.png', img)
51 |
52 | @staticmethod
53 | def get_styles():
54 | return ['intensive', 'less intensive', 'normal', 'soft']
55 |
56 |
57 | class ICLightBeforeUpsacle(ICLight):
58 |
59 | def preprocess(self, context: Context, image: Image):
60 | super().preprocess(context, image)
61 | return self.enabled and self.enable_before_upscale
62 |
63 |
64 | class ICLightAfterUpsacle(ICLight):
65 |
66 | def preprocess(self, context: Context, image: Image):
67 | super().preprocess(context, image)
68 | return self.enabled and not self.enable_before_upscale
69 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/img2imgmasking.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from modules import devices
4 | from modules.processing import StableDiffusionProcessingImg2Img
5 |
6 | from sd_bmab import masking
7 | from sd_bmab import external
8 | from sd_bmab.base.context import Context
9 | from sd_bmab.base.processorbase import ProcessorBase
10 |
11 |
12 | class Img2imgMasking(ProcessorBase):
13 | def __init__(self) -> None:
14 | super().__init__()
15 | self.enabled = False
16 | self.prompt = ''
17 | self.input_image = None
18 |
19 | def preprocess(self, context: Context, image: Image):
20 | self.enabled = context.args['detect_enabled']
21 | self.prompt = context.args['masking_prompt']
22 | self.input_image = context.args['input_image']
23 | return not context.is_txtimg() and self.enabled
24 |
25 | def sam(self, context, prompt, input_image):
26 | with external.ModuleAutoLoader('groundingdino', 'grdino') as dino:
27 | boxes, logits, phrases = dino.dino_predict(input_image, prompt, 0.35, 0.25)
28 | sam = masking.get_mask_generator()
29 | mask = sam.predict(input_image, boxes)
30 | return mask
31 |
32 | def process(self, context: Context, image: Image):
33 | if context.sdprocessing.image_mask is not None:
34 | context.sdprocessing.image_mask = self.sam(self.prompt, context.sdprocessing.init_images[0])
35 | context.script.extra_image.append(context.sdprocessing.image_mask)
36 | if context.sdprocessing.image_mask is None and self.input_image is not None:
37 | mask = self.sam(context, self.prompt, context.sdprocessing.init_images[0])
38 | newpil = Image.new('RGB', context.sdprocessing.init_images[0].size)
39 | newdata = [bdata if mdata == 0 else ndata for mdata, ndata, bdata in
40 | zip(mask.getdata(), context.sdprocessing.init_images[0].getdata(), self.input_image.getdata())]
41 | newpil.putdata(newdata)
42 | context.script.extra_image.append(newpil)
43 | return newpil
44 | return image
45 |
46 | def postprocess(self, context: Context, image: Image):
47 | devices.torch_gc()
48 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/intermediate.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from sd_bmab import util
3 | from sd_bmab.base.context import Context
4 | from sd_bmab.base.processorbase import ProcessorBase
5 |
6 |
7 | class NoiseAlpha(ProcessorBase):
8 | def __init__(self) -> None:
9 | super().__init__()
10 | self.noise_alpha = 0
11 |
12 | def preprocess(self, context: Context, image: Image):
13 | self.noise_alpha = context.args['noise_alpha']
14 | return self.noise_alpha != 0
15 |
16 | def process(self, context: Context, image: Image):
17 | context.add_generation_param('BMAB noise alpha', self.noise_alpha)
18 | img_noise = util.generate_noise(context.sdprocessing.seed, image.size[0], image.size[1])
19 | return Image.blend(image, img_noise, alpha=self.noise_alpha)
20 |
21 | def postprocess(self, context: Context, image: Image):
22 | pass
23 |
--------------------------------------------------------------------------------
/sd_bmab/processors/basic/preprocessfilter.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.base import filter
2 |
3 |
4 | def run_preprocess_filter(context):
5 | module_config = context.args.get('module_config', {})
6 | filter_name = module_config.get('preprocess_filter', None)
7 | if filter_name is None or filter_name == 'None':
8 | return
9 |
10 | bmab_filter = filter.get_filter(filter_name)
11 | filter.preprocess_filter(bmab_filter, context, None)
12 | filter.process_filter(bmab_filter, context, None, None)
13 | filter.postprocess_filter(bmab_filter, context)
14 |
--------------------------------------------------------------------------------
/sd_bmab/processors/controlnet/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.controlnet.noise import LineartNoise
2 | from sd_bmab.processors.controlnet.pose import Openpose
3 | from sd_bmab.processors.controlnet.ipadapter import IpAdapter
4 |
--------------------------------------------------------------------------------
/sd_bmab/processors/controlnet/noise.py:
--------------------------------------------------------------------------------
1 | import os
2 | from PIL import Image
3 |
4 | from modules import shared
5 |
6 | import sd_bmab
7 | from sd_bmab import util, controlnet
8 | from sd_bmab.util import debug_print
9 | from sd_bmab.base.context import Context
10 | from sd_bmab.base.processorbase import ProcessorBase
11 | from sd_bmab.base import cache
12 |
13 |
14 | class LineartNoise(ProcessorBase):
15 | def __init__(self) -> None:
16 | super().__init__()
17 | self.controlnet_opt = {}
18 | self.enabled = False
19 | self.with_refiner = False
20 | self.noise_enabled = False
21 | self.noise_strength = 0.4
22 | self.noise_begin = 0.1
23 | self.noise_end = 0.9
24 | self.noise_hiresfix = 'Both'
25 |
26 | @staticmethod
27 | def with_refiner(context: Context):
28 | controlnet_opt = context.args.get('module_config', {}).get('controlnet', {})
29 | enabled = controlnet_opt.get('enabled', False)
30 | with_refiner = controlnet_opt.get('with_refiner', False)
31 | debug_print('with refiner', enabled, with_refiner)
32 | return enabled and with_refiner
33 |
34 | def preprocess(self, context: Context, image: Image):
35 | self.controlnet_opt = context.args.get('module_config', {}).get('controlnet', {})
36 | self.enabled = self.controlnet_opt.get('enabled', False)
37 | self.with_refiner = self.controlnet_opt.get('with_refiner', False)
38 | self.noise_enabled = self.controlnet_opt.get('noise', False)
39 | self.noise_strength = self.controlnet_opt.get('noise_strength', 0.4)
40 | self.noise_begin = self.controlnet_opt.get('noise_begin', 0.1)
41 | self.noise_end = self.controlnet_opt.get('noise_end', 0.9)
42 | self.noise_hiresfix = self.controlnet_opt.get('noise_hiresfix', self.noise_hiresfix)
43 |
44 | debug_print('Noise', context.is_refiner_context(), context.with_refiner(), self.with_refiner)
45 | if context.is_refiner_context():
46 | return self.enabled and self.noise_enabled and self.with_refiner
47 | elif context.with_refiner() and self.with_refiner:
48 | return False
49 | return self.enabled and self.noise_enabled
50 |
51 | @staticmethod
52 | def get_noise_args(image, weight, begin, end, hr_option):
53 | cn_args = {
54 | 'enabled': True,
55 | 'image': image if isinstance(image, str) and os.path.exists(image) else util.b64_encoding(image.convert('RGB')),
56 | 'model': shared.opts.bmab_cn_lineart,
57 | 'weight': weight,
58 | "guidance_start": begin,
59 | "guidance_end": end,
60 | 'resize_mode': 'Just Resize',
61 | 'pixel_perfect': False,
62 | 'control_mode': 'ControlNet is more important',
63 | 'processor_res': 512,
64 | 'threshold_a': 0.5,
65 | 'threshold_b': 0.5,
66 | 'hr_option': hr_option
67 | }
68 | return cn_args
69 |
70 | def get_controlnet_args(self, context):
71 | img = util.generate_noise(context.sdprocessing.seed, context.sdprocessing.width, context.sdprocessing.height)
72 | noise = img.convert('L').convert('RGB')
73 | return self.get_noise_args(noise, self.noise_strength, self.noise_begin, self.noise_end, self.noise_hiresfix)
74 |
75 | def process(self, context: Context, image: Image):
76 | context.add_generation_param('BMAB controlnet mode', 'lineart')
77 | context.add_generation_param('BMAB noise strength', self.noise_strength)
78 | context.add_generation_param('BMAB noise begin', self.noise_begin)
79 | context.add_generation_param('BMAB noise end', self.noise_end)
80 |
81 | index = controlnet.get_controlnet_index(context.sdprocessing)
82 | img = cache.get_noise_from_cache(context.sdprocessing.seed, context.sdprocessing.width, context.sdprocessing.height)
83 | cn_op_arg = self.get_noise_args(img, self.noise_strength, self.noise_begin, self.noise_end, self.noise_hiresfix)
84 | debug_print(f'Noise Enabled {index}')
85 | sc_args = list(context.sdprocessing.script_args)
86 | sc_args[index] = cn_op_arg
87 | context.sdprocessing.script_args = tuple(sc_args)
88 |
89 | def postprocess(self, context: Context, image: Image):
90 | pass
91 |
92 |
93 |
--------------------------------------------------------------------------------
/sd_bmab/processors/controlnet/pose.py:
--------------------------------------------------------------------------------
1 | import os
2 | import glob
3 | import random
4 |
5 | from PIL import Image
6 |
7 | from modules import shared
8 |
9 | import sd_bmab
10 | from sd_bmab import util, controlnet
11 | from sd_bmab.util import debug_print
12 | from sd_bmab.base.context import Context
13 | from sd_bmab.base.processorbase import ProcessorBase
14 |
15 |
16 | class Openpose(ProcessorBase):
17 | def __init__(self) -> None:
18 | super().__init__()
19 | self.controlnet_opt = {}
20 | self.enabled = False
21 | self.pose_enabled = False
22 | self.pose_strength = 0.3
23 | self.pose_begin = 0.0
24 | self.pose_end = 1.0
25 | self.pose_face_only = False
26 | self.pose_selected = 'Random'
27 |
28 | def preprocess(self, context: Context, image: Image):
29 | self.controlnet_opt = context.args.get('module_config', {}).get('controlnet', {})
30 | self.enabled = self.controlnet_opt.get('enabled', False)
31 | self.pose_enabled = self.controlnet_opt.get('pose', False)
32 | self.pose_strength = self.controlnet_opt.get('pose_strength', self.pose_strength)
33 | self.pose_begin = self.controlnet_opt.get('pose_begin', self.pose_begin)
34 | self.pose_end = self.controlnet_opt.get('pose_end', self.pose_end)
35 | self.pose_face_only = self.controlnet_opt.get('pose_face_only', self.pose_face_only)
36 | self.pose_selected = self.controlnet_opt.get('pose_selected', self.pose_selected)
37 | return self.enabled and self.pose_enabled
38 |
39 | def get_openpose_args(self, image):
40 | cn_args = {
41 | 'enabled': True,
42 | 'image': image if isinstance(image, str) and os.path.exists(image) else util.b64_encoding(image.convert('RGB')),
43 | 'module': 'openpose_faceonly' if self.pose_face_only else 'openpose',
44 | 'model': shared.opts.bmab_cn_openpose,
45 | 'weight': self.pose_strength,
46 | "guidance_start": self.pose_begin,
47 | "guidance_end": self.pose_end,
48 | 'resize_mode': 'Just Resize',
49 | 'pixel_perfect': False,
50 | 'control_mode': 'My prompt is more important',
51 | 'processor_res': 512,
52 | 'threshold_a': 0.5,
53 | 'threshold_b': 0.5,
54 | 'hr_option': 'Low res only'
55 | }
56 | return cn_args
57 |
58 | def process(self, context: Context, image: Image):
59 | context.add_generation_param('BMAB controlnet pose mode', 'openpose')
60 | context.add_generation_param('BMAB pose strength', self.pose_strength)
61 | context.add_generation_param('BMAB pose begin', self.pose_begin)
62 | context.add_generation_param('BMAB pose end', self.pose_end)
63 |
64 | img = self.load_random_image(context)
65 | if img is None:
66 | return
67 |
68 | index = controlnet.get_controlnet_index(context.sdprocessing)
69 | cn_op_arg = self.get_openpose_args(img)
70 | debug_print(f'Pose Enabled {index}')
71 | sc_args = list(context.sdprocessing.script_args)
72 | sc_args[index] = cn_op_arg
73 | context.sdprocessing.script_args = tuple(sc_args)
74 |
75 | def postprocess(self, context: Context, image: Image):
76 | pass
77 |
78 | def load_random_image(self, context):
79 | path = os.path.dirname(sd_bmab.__file__)
80 | path = os.path.normpath(os.path.join(path, '../resources/pose'))
81 | if os.path.exists(path) and os.path.isdir(path):
82 | file_mask = f'{path}/*.*'
83 | files = glob.glob(file_mask)
84 | if not files:
85 | debug_print(f'Not found pose files in {path}')
86 | return None
87 | img = context.load('preprocess_image')
88 | if img is not None:
89 | return img
90 | if self.pose_selected == 'Random':
91 | file = random.choice(files)
92 | debug_print(f'Random pose {file}')
93 | return self.get_cache(context, file)
94 | else:
95 | img_name = f'{path}/{self.pose_selected}'
96 | return self.get_cache(context, img_name)
97 | debug_print(f'Not found directory {path}')
98 | return None
99 |
100 | def get_cache(self, context, file):
101 | if self.pose_face_only:
102 | path = os.path.dirname(sd_bmab.__file__)
103 | path = os.path.normpath(os.path.join(path, '../resources/cache'))
104 | if os.path.exists(path) and os.path.isdir(path):
105 | b = os.path.basename(file)
106 | file_mask = f'{path}/pose_face_{b}'
107 | if os.path.exists(file_mask):
108 | return Image.open(file_mask)
109 | return Image.open(file)
110 |
111 |
112 | @staticmethod
113 | def list_pose():
114 | path = os.path.dirname(sd_bmab.__file__)
115 | path = os.path.normpath(os.path.join(path, '../resources/pose'))
116 | if os.path.exists(path) and os.path.isdir(path):
117 | file_mask = f'{path}/*.*'
118 | files = glob.glob(file_mask)
119 | return [os.path.basename(f) for f in files]
120 | debug_print(f'Not found directory {path}')
121 | return []
122 |
123 | @staticmethod
124 | def get_pose(f):
125 | if f == 'Random' or f == 'Preprocess':
126 | return Image.new('RGB', (512, 512), 0)
127 | path = os.path.dirname(sd_bmab.__file__)
128 | path = os.path.normpath(os.path.join(path, '../resources/pose'))
129 | if os.path.exists(path) and os.path.isdir(path):
130 | img_name = f'{path}/{f}'
131 | return Image.open(img_name)
132 | return Image.new('RGB', (512, 512), 0)
133 |
134 | @staticmethod
135 | def pose_selected(*args):
136 | return Openpose.get_pose(args[0])
137 |
--------------------------------------------------------------------------------
/sd_bmab/processors/detailer/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.detailer.face import FaceDetailer, PreprocessFaceDetailer, FaceDetailerBeforeUpsacle
2 | from sd_bmab.processors.detailer.hand import HandDetailer
3 | from sd_bmab.processors.detailer.person import PersonDetailer
4 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.postprocess.upscaleafterprocess import AfterProcessUpscaler
2 | from sd_bmab.processors.postprocess.upscalebeforeprocess import BeforeProcessUpscaler
3 | from sd_bmab.processors.postprocess.inpaint import InpaintResize
4 | from sd_bmab.processors.postprocess.inpaintlama import InpaintLamaResize
5 | from sd_bmab.processors.postprocess.finalfilter import FinalFilter
6 | from sd_bmab.processors.postprocess.watermark import Watermark
7 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/finalfilter.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from sd_bmab.base import Context, ProcessorBase
3 | from sd_bmab.base import filter
4 |
5 |
6 | class FinalFilter(ProcessorBase):
7 | def __init__(self) -> None:
8 | super().__init__()
9 | self.filter_name = 'None'
10 | self.filter = None
11 |
12 | def preprocess(self, context: Context, image: Image):
13 | self.filter_name = context.args.get('postprocess_final_filter', self.filter_name)
14 | return self.filter_name != 'None'
15 |
16 | def process(self, context: Context, image: Image):
17 | self.filter = filter.get_filter(self.filter_name)
18 | if self.filter is not None:
19 | filter.preprocess_filter(self.filter, context, image)
20 | image = filter.process_filter(self.filter, context, None, image)
21 | filter.postprocess_filter(self.filter, context)
22 | return image
23 |
24 | def finalprocess(self, context: Context, image: Image):
25 | if self.filter is not None:
26 | filter.finalprocess_filter(self.filter, context)
27 | self.filter = None
28 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/inpaint.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from PIL import ImageDraw
3 |
4 | from modules import devices
5 | from modules.processing import process_images, StableDiffusionProcessingImg2Img
6 |
7 | from sd_bmab import util
8 | from sd_bmab.base import apply_extensions, build_img2img, Context, ProcessorBase
9 |
10 | from sd_bmab.util import debug_print
11 | from sd_bmab.detectors import UltralyticsPersonDetector8n
12 |
13 |
14 | class InpaintResize(ProcessorBase):
15 | def __init__(self) -> None:
16 | super().__init__()
17 | self.resize_by_person_opt = None
18 | self.value = 0,
19 | self.denoising_strength = 0.4
20 | self.dilation = 4
21 | self.mode = 'Inpaint'
22 |
23 | def preprocess(self, context: Context, image: Image):
24 | enabled = context.args.get('resize_by_person_enabled', False)
25 | self.resize_by_person_opt = context.args.get('module_config', {}).get('resize_by_person_opt', {})
26 | self.value = self.resize_by_person_opt.get('scale', 0)
27 | self.denoising_strength = self.resize_by_person_opt.get('denoising_strength', 0.4)
28 | self.dilation = self.resize_by_person_opt.get('dilation', 0.4)
29 | self.mode = self.resize_by_person_opt.get('mode', self.mode)
30 |
31 | return enabled and self.mode == 'Inpaint'
32 |
33 | def process(self, context: Context, image: Image):
34 | debug_print('prepare detector')
35 | detector = UltralyticsPersonDetector8n()
36 | boxes, logits = detector.predict(context, image)
37 |
38 | org_size = image.size
39 | debug_print('size', org_size)
40 |
41 | largest = (0, None)
42 | for box in boxes:
43 | x1, y1, x2, y2 = box
44 | size = (x2 - x1) * (y2 - y1)
45 | if size > largest[0]:
46 | largest = (size, box)
47 |
48 | if largest[0] == 0:
49 | return image
50 |
51 | x1, y1, x2, y2 = largest[1]
52 | ratio = (y2 - y1) / image.height
53 | debug_print('ratio', ratio)
54 | debug_print('org_size', org_size)
55 | if ratio <= self.value:
56 | return image
57 | image_ratio = ratio / self.value
58 | if image_ratio < 1.0:
59 | return image
60 | debug_print('scale', image_ratio)
61 | ratio = image_ratio
62 |
63 | org_size = image.size
64 | dw, dh = org_size
65 |
66 | context.add_generation_param('BMAB controlnet mode', 'inpaint')
67 | context.add_generation_param('BMAB resize by person ratio', '%.3s' % ratio)
68 |
69 | resized_width = int(dw / ratio)
70 | resized_height = int(dh / ratio)
71 | resized = image.resize((resized_width, resized_height), resample=util.LANCZOS)
72 | context.sdprocessing.resize_mode = 2
73 | input_image = util.resize_image(2, resized, dw, dh)
74 |
75 | offset_x = int((dw - resized_width) / 2)
76 | offset_y = dh - resized_height
77 |
78 | mask = Image.new('L', (dw, dh), 255)
79 | dr = ImageDraw.Draw(mask, 'L')
80 | dr.rectangle((offset_x, offset_y, offset_x + resized_width, offset_y + resized_height), fill=0)
81 | mask = mask.resize(org_size, resample=util.LANCZOS)
82 | mask = util.dilate_mask(mask, self.dilation)
83 |
84 | opt = dict(mask=mask, denoising_strength=self.denoising_strength)
85 | i2i_param = build_img2img(context.sdprocessing, input_image, opt)
86 |
87 | img2img = StableDiffusionProcessingImg2Img(**i2i_param)
88 | img2img.cached_c = [None, None]
89 | img2img.cached_uc = [None, None]
90 | img2img.scripts, img2img.script_args = apply_extensions(context.sdprocessing, cn_enabled=False)
91 |
92 | processed = process_images(img2img)
93 | img = processed.images[0]
94 |
95 | img2img.close()
96 |
97 | return img
98 |
99 | def postprocess(self, context: Context, image: Image):
100 | devices.torch_gc()
101 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/inpaintlama.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from PIL import ImageDraw
3 |
4 | from modules import shared
5 | from modules import devices
6 | from modules.processing import process_images, StableDiffusionProcessingImg2Img
7 |
8 | from sd_bmab import util
9 | from sd_bmab.base import apply_extensions, build_img2img, Context, ProcessorBase
10 | from sd_bmab.util import debug_print
11 | from sd_bmab.detectors import UltralyticsPersonDetector8n
12 |
13 |
14 | class InpaintLamaResize(ProcessorBase):
15 | def __init__(self) -> None:
16 | super().__init__()
17 | self.resize_by_person_opt = None
18 | self.value = 0,
19 | self.denoising_strength = 0.4
20 | self.dilation = 4
21 | self.mode = 'ControlNet inpaint+lama'
22 | self.enabled = False
23 |
24 | def use_controlnet(self, context: Context):
25 | self.preprocess(context, None)
26 | return self.enabled
27 |
28 | def preprocess(self, context: Context, image: Image):
29 | self.enabled = context.args.get('resize_by_person_enabled', self.enabled)
30 | self.resize_by_person_opt = context.args.get('module_config', {}).get('resize_by_person_opt', {})
31 | self.value = self.resize_by_person_opt.get('scale', 0)
32 | self.denoising_strength = self.resize_by_person_opt.get('denoising_strength', 0.4)
33 | self.dilation = self.resize_by_person_opt.get('dilation', 0.4)
34 | self.mode = self.resize_by_person_opt.get('mode', self.mode)
35 |
36 | return self.enabled and self.mode == 'ControlNet inpaint+lama'
37 |
38 | @staticmethod
39 | def get_inpaint_lama_args(image, mask):
40 | cn_args = {
41 | 'input_image': util.b64_encoding(image),
42 | 'mask': util.b64_encoding(mask),
43 | 'module': 'inpaint_only+lama',
44 | 'model': shared.opts.bmab_cn_inpaint,
45 | 'weight': 1,
46 | "guidance_start": 0,
47 | "guidance_end": 1,
48 | 'resize_mode': 'Resize and Fill',
49 | 'pixel_perfect': False,
50 | 'control_mode': 'ControlNet is more important',
51 | 'processor_res': 512,
52 | 'threshold_a': 64,
53 | 'threshold_b': 64,
54 | }
55 | return cn_args
56 |
57 | def get_ratio(self, context, img, p):
58 | p.extra_generation_params['BMAB process_resize_by_person'] = self.value
59 |
60 | final_ratio = 1
61 | debug_print('prepare detector')
62 | detector = UltralyticsPersonDetector8n()
63 | boxes, logits = detector.predict(context, img)
64 |
65 | largest = (0, None)
66 | for box in boxes:
67 | x1, y1, x2, y2 = box
68 | size = (x2 - x1) * (y2 - y1)
69 | if size > largest[0]:
70 | largest = (size, box)
71 |
72 | if largest[0] == 0:
73 | return final_ratio
74 |
75 | x1, y1, x2, y2 = largest[1]
76 | ratio = (y2 - y1) / img.height
77 | debug_print('ratio', ratio)
78 |
79 | if ratio > self.value:
80 | image_ratio = ratio / self.value
81 | if image_ratio < 1.0:
82 | return final_ratio
83 | final_ratio = image_ratio
84 | return final_ratio
85 |
86 | def resize_by_person_using_controlnet(self, context, p):
87 | if not isinstance(p, StableDiffusionProcessingImg2Img):
88 | return False
89 |
90 | cn_args = util.get_cn_args(p)
91 |
92 | debug_print('resize_by_person_enabled_inpaint', self.value)
93 | img = p.init_images[0]
94 | context.script.extra_image.append(img)
95 |
96 | ratio = self.get_ratio(context, img, p)
97 | debug_print('image resize ratio', ratio)
98 | org_size = img.size
99 | dw, dh = org_size
100 |
101 | if ratio == 1:
102 | return False
103 |
104 | p.extra_generation_params['BMAB controlnet mode'] = 'inpaint+lama'
105 | p.extra_generation_params['BMAB resize by person ratio'] = '%.3s' % ratio
106 |
107 | resized_width = int(dw / ratio)
108 | resized_height = int(dh / ratio)
109 | resized = img.resize((resized_width, resized_height), resample=util.LANCZOS)
110 | p.resize_mode = 2
111 | input_image = util.resize_image(2, resized, dw, dh)
112 | p.init_images[0] = input_image
113 |
114 | offset_x = int((dw - resized_width) / 2)
115 | offset_y = dh - resized_height
116 |
117 | mask = Image.new('L', (dw, dh), 255)
118 | dr = ImageDraw.Draw(mask, 'L')
119 | dr.rectangle((offset_x, offset_y, offset_x + resized_width, offset_y + resized_height), fill=0)
120 | mask = mask.resize(org_size, resample=util.LANCZOS)
121 | mask = util.dilate_mask(mask, self.dilation)
122 |
123 | cn_op_arg = self.get_inpaint_lama_args(input_image, mask)
124 | idx = cn_args[0]
125 | sc_args = list(p.script_args)
126 | sc_args[idx] = cn_op_arg
127 | p.script_args = tuple(sc_args)
128 | return True
129 |
130 | def process(self, context: Context, image: Image):
131 | opt = dict(denoising_strength=self.denoising_strength)
132 | i2i_param = build_img2img(context.sdprocessing, image, opt)
133 |
134 | img2img = StableDiffusionProcessingImg2Img(**i2i_param)
135 | img2img.cached_c = [None, None]
136 | img2img.cached_uc = [None, None]
137 | img2img.scripts, img2img.script_args = apply_extensions(context.sdprocessing, cn_enabled=True)
138 |
139 | if self.resize_by_person_using_controlnet(context, img2img):
140 | processed = process_images(img2img)
141 | image = processed.images[0]
142 | img2img.close()
143 | devices.torch_gc()
144 | return image
145 |
146 | def postprocess(self, context: Context, image: Image):
147 | pass
148 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/upscaleafterprocess.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from modules import images
4 |
5 | from sd_bmab.base.context import Context
6 | from sd_bmab.base.processorbase import ProcessorBase
7 | from sd_bmab.util import debug_print
8 |
9 |
10 | class AfterProcessUpscaler(ProcessorBase):
11 | def __init__(self) -> None:
12 | super().__init__()
13 | self.ratio = 1.5
14 | self.upscaler = 'None'
15 |
16 | def preprocess(self, context: Context, image: Image):
17 | self.ratio = context.args['upscale_ratio']
18 | self.upscaler = context.args['upscaler_name']
19 | return context.args['upscale_enabled'] and context.args['detailing_after_upscale']
20 |
21 | def process(self, context: Context, image: Image):
22 | debug_print(f'Upscale ratio {self.ratio} Upscaler {self.upscaler}')
23 | context.add_generation_param('BMAB_upscale_option', f'Upscale ratio {self.ratio} Upscaler {self.upscaler}')
24 |
25 | if self.ratio < 1.0 or self.ratio > 4.0:
26 | debug_print('upscale out of range')
27 | return image
28 | image = image.convert('RGB')
29 | context.add_generation_param('BMAB process upscale', self.ratio)
30 | context.args['max_area'] = image.width * image.height
31 | context.args['upscale_limit'] = True
32 |
33 | w = image.width
34 | h = image.height
35 | img = images.resize_image(0, image, int(w * self.ratio), int(h * self.ratio), self.upscaler)
36 | return img.convert('RGB')
37 |
38 | def postprocess(self, context: Context, image: Image):
39 | pass
40 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/upscalebeforeprocess.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from modules import images
4 |
5 | from sd_bmab.base.context import Context
6 | from sd_bmab.base.processorbase import ProcessorBase
7 | from sd_bmab.util import debug_print
8 |
9 |
10 | class BeforeProcessUpscaler(ProcessorBase):
11 | def __init__(self) -> None:
12 | super().__init__()
13 | self.ratio = 1.5
14 | self.upscaler = 'None'
15 |
16 | def preprocess(self, context: Context, image: Image):
17 | self.ratio = context.args['upscale_ratio']
18 | self.upscaler = context.args['upscaler_name']
19 | return context.args['upscale_enabled'] and not context.args['detailing_after_upscale']
20 |
21 | def process(self, context: Context, image: Image):
22 | debug_print(f'Upscale ratio {self.ratio} Upscaler {self.upscaler}')
23 | context.add_generation_param('BMAB_upscale_option', f'Upscale ratio {self.ratio} Upscaler {self.upscaler}')
24 |
25 | if self.ratio < 1.0 or self.ratio > 4.0:
26 | debug_print('upscale out of range')
27 | return image
28 | image = image.convert('RGB')
29 | context.add_generation_param('BMAB process upscale', self.ratio)
30 | context.args['upscale_limit'] = True
31 |
32 | w = image.width
33 | h = image.height
34 | img = images.resize_image(0, image, int(w * self.ratio), int(h * self.ratio), self.upscaler)
35 | return img.convert('RGB')
36 |
37 | def postprocess(self, context: Context, image: Image):
38 | pass
39 |
--------------------------------------------------------------------------------
/sd_bmab/processors/postprocess/watermark.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import glob
4 |
5 | from PIL import Image
6 | from PIL import ImageDraw
7 | from PIL import ImageFont
8 | from sd_bmab.base import Context, ProcessorBase
9 |
10 |
11 | class Watermark(ProcessorBase):
12 | alignment = {
13 | 'top': lambda w, h, cx, cy: (w / 2 - cx / 2, 0),
14 | 'top-right': lambda w, h, cx, cy: (w - cx, 0),
15 | 'right': lambda w, h, cx, cy: (w - cx, h / 2 - cy / 2),
16 | 'bottom-right': lambda w, h, cx, cy: (w - cx, h - cy),
17 | 'bottom': lambda w, h, cx, cy: (w / 2 - cx / 2, h - cy),
18 | 'bottom-left': lambda w, h, cx, cy: (0, h - cy),
19 | 'left': lambda w, h, cx, cy: (0, h / 2 - cy / 2),
20 | 'top-left': lambda w, h, cx, cy: (0, 0),
21 | 'center': lambda w, h, cx, cy: (w / 2 - cx / 2, h / 2 - cy / 2),
22 | }
23 |
24 | def __init__(self) -> None:
25 | super().__init__()
26 | self.enabled = False
27 | self.font = None
28 | self.alignment = 'bottom-left'
29 | self.text_alignment = 'left'
30 | self.rotate = 0
31 | self.color = '#000000'
32 | self.background_color = '#00000000'
33 | self.font_size = 12
34 | self.transparency = 100
35 | self.background_transparency = 0
36 | self.margin = 5
37 | self.text = ''
38 |
39 | def preprocess(self, context: Context, image: Image):
40 | watermark_opt = context.args.get('module_config', {}).get('watermark', {})
41 | self.enabled = watermark_opt.get('enabled', self.enabled)
42 | self.font = watermark_opt.get('font', self.font)
43 | self.alignment = watermark_opt.get('alignment', self.alignment)
44 | self.text_alignment = watermark_opt.get('text_alignment', self.text_alignment)
45 | _rotate = watermark_opt.get('rotate', self.rotate)
46 | self.rotate = int(_rotate)
47 | self.color = watermark_opt.get('color', self.color)
48 | self.background_color = watermark_opt.get('background_color', self.background_color)
49 | self.font_size = watermark_opt.get('font_size', self.font_size)
50 | self.transparency = watermark_opt.get('transparency', self.transparency)
51 | self.background_transparency = watermark_opt.get('background_transparency', self.background_transparency)
52 | self.margin = watermark_opt.get('margin', self.margin)
53 | self.text = watermark_opt.get('text', self.text)
54 |
55 | return self.enabled
56 |
57 | def process(self, context: Context, image: Image):
58 |
59 | background_color = self.color_hex_to_rgb(self.background_color, int(255 * (self.background_transparency / 100)))
60 |
61 | if os.path.isfile(self.text):
62 | cropped = Image.open(self.text)
63 | else:
64 | font = self.get_font(self.font, self.font_size)
65 | color = self.color_hex_to_rgb(self.color, int(255 * (self.transparency / 100)))
66 |
67 | # 1st
68 | base = Image.new('RGBA', image.size, background_color)
69 | draw = ImageDraw.Draw(base)
70 | bbox = draw.textbbox((0, 0), self.text, font=font)
71 | draw.text((0, 0), self.text, font=font, fill=color, align=self.text_alignment)
72 | cropped = base.crop(bbox)
73 |
74 | # 2st margin
75 | base = Image.new('RGBA', (cropped.width + self.margin * 2, cropped.height + self.margin * 2), background_color)
76 | base.paste(cropped, (self.margin, self.margin))
77 |
78 | # 3rd rotate
79 | base = base.rotate(self.rotate, expand=True)
80 |
81 | # 4th
82 | image = image.convert('RGBA')
83 | image2 = image.copy()
84 | x, y = Watermark.alignment[self.alignment](image.width, image.height, base.width, base.height)
85 | image2.paste(base, (int(x), int(y)))
86 | image = Image.alpha_composite(image, image2)
87 |
88 | return image.convert('RGB')
89 |
90 | @staticmethod
91 | def color_hex_to_rgb(value, transparency):
92 | value = value.lstrip('#')
93 | lv = len(value)
94 | r, g, b = tuple(int(value[i:i + 2], 16) for i in range(0, lv, 2))
95 | return r, g, b, transparency
96 |
97 | @staticmethod
98 | def list_fonts():
99 | if sys.platform == 'win32':
100 | path = 'C:\\Windows\\Fonts\\*.ttf'
101 | files = glob.glob(path)
102 | return [os.path.basename(f) for f in files]
103 | if sys.platform == 'darwin':
104 | path = '/System/Library/Fonts/*'
105 | files = glob.glob(path)
106 | return [os.path.basename(f) for f in files]
107 | if sys.platform == 'linux':
108 | path = '/usr/share/fonts/*'
109 | files = glob.glob(path)
110 | fonts = [os.path.basename(f) for f in files]
111 | if 'SAGEMAKER_INTERNAL_IMAGE_URI' in os.environ:
112 | path = '/opt/conda/envs/sagemaker-distribution/fonts/*'
113 | files = glob.glob(path)
114 | fonts.extend([os.path.basename(f) for f in files])
115 | return fonts
116 | return ['']
117 |
118 | @staticmethod
119 | def get_font(font, size):
120 | if sys.platform == 'win32':
121 | path = f'C:\\Windows\\Fonts\\{font}'
122 | return ImageFont.truetype(path, size, encoding="unic")
123 | if sys.platform == 'darwin':
124 | path = f'/System/Library/Fonts/{font}'
125 | return ImageFont.truetype(path, size, encoding="unic")
126 | if sys.platform == 'linux':
127 | if 'SAGEMAKER_INTERNAL_IMAGE_URI' in os.environ:
128 | path = f'/opt/conda/envs/sagemaker-distribution/fonts/{font}'
129 | else:
130 | path = f'/usr/share/fonts/{font}'
131 | return ImageFont.truetype(path, size, encoding="unic")
132 |
--------------------------------------------------------------------------------
/sd_bmab/processors/preprocess/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.preprocess.resize import ResizeIntermidiateBeforeUpscale, ResizeIntermidiateAfterUpsacle
2 | from sd_bmab.processors.preprocess.refiner import RefinerPreprocessor
3 | from sd_bmab.processors.preprocess.pretraining import PretrainingDetailer, PretrainingDetailerBeforeUpscale
4 | from sd_bmab.processors.preprocess.resample import ResamplePreprocessor, ResamplePreprocessorBeforeUpscale
5 | from sd_bmab.processors.preprocess.preprocessfilter import PreprocessFilter
6 |
7 |
--------------------------------------------------------------------------------
/sd_bmab/processors/preprocess/preprocessfilter.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 | from sd_bmab.base import Context, ProcessorBase
3 | from sd_bmab.base import filter
4 |
5 |
6 | class PreprocessFilter(ProcessorBase):
7 | def __init__(self) -> None:
8 | super().__init__()
9 | self.filter_name = 'None'
10 | self.filter = None
11 |
12 | def preprocess(self, context: Context, image: Image):
13 | self.filter_name = context.args.get('txt2img_filter_hresfix_before_upscale', self.filter_name)
14 | return not context.is_hires_fix() and self.filter_name != 'None'
15 |
16 | def process(self, context: Context, image: Image):
17 | self.filter = filter.get_filter(self.filter_name)
18 | if self.filter is not None:
19 | filter.preprocess_filter(self.filter, context, image)
20 | image = filter.process_filter(self.filter, context, None, image)
21 | filter.postprocess_filter(self.filter, context)
22 | return image
23 |
24 | def finalprocess(self, context: Context, image: Image):
25 | if self.filter is not None:
26 | filter.finalprocess_filter(self.filter, context)
27 | self.filter = None
28 |
--------------------------------------------------------------------------------
/sd_bmab/processors/preprocess/refiner.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from modules import devices
4 | from modules import images
5 |
6 | from sd_bmab import util
7 | from sd_bmab import constants
8 | from sd_bmab.util import debug_print
9 | from sd_bmab.base import process_img2img, Context, ProcessorBase, process_img2img_with_controlnet
10 | from sd_bmab.processors.controlnet import LineartNoise
11 |
12 |
13 | class RefinerPreprocessor(ProcessorBase):
14 | def __init__(self) -> None:
15 | super().__init__()
16 |
17 | self.refiner_opt = {}
18 | self.enabled = False
19 | self.checkpoint = constants.checkpoint_default
20 | self.vae = constants.vae_default
21 | self.keep_checkpoint = True
22 | self.prompt = None
23 | self.negative_prompt = None
24 | self.sampler = None
25 | self.scheduler = None
26 | self.upscaler = None
27 | self.steps = 20
28 | self.cfg_scale = 0.7
29 | self.denoising_strength = 0.75
30 | self.scale = 1
31 | self.width = 0
32 | self.height = 0
33 |
34 | def preprocess(self, context: Context, image: Image):
35 | self.enabled = context.args['refiner_enabled']
36 | self.refiner_opt = context.args.get('module_config', {}).get('refiner_opt', {})
37 |
38 | self.checkpoint = self.refiner_opt.get('checkpoint', self.checkpoint)
39 | self.vae = self.refiner_opt.get('vae', self.vae)
40 | self.keep_checkpoint = self.refiner_opt.get('keep_checkpoint', True)
41 | self.prompt = self.refiner_opt.get('prompt', '')
42 | self.negative_prompt = self.refiner_opt.get('negative_prompt', '')
43 | self.sampler = self.refiner_opt.get('sampler', None)
44 | self.scheduler = self.refiner_opt.get('scheduler', None)
45 | self.upscaler = self.refiner_opt.get('upscaler', None)
46 | self.steps = self.refiner_opt.get('steps', None)
47 | self.cfg_scale = self.refiner_opt.get('cfg_scale', None)
48 | self.denoising_strength = self.refiner_opt.get('denoising_strength', None)
49 | self.scale = self.refiner_opt.get('scale', None)
50 | self.width = self.refiner_opt.get('width', None)
51 | self.height = self.refiner_opt.get('height', None)
52 |
53 | if self.enabled:
54 | context.refiner = self
55 |
56 | return self.enabled
57 |
58 | def process(self, context: Context, image: Image):
59 | output_width = image.width
60 | output_height = image.height
61 |
62 | if not (self.width == 0 and self.height == 0 and self.scale == 1):
63 | if (self.width == 0 or self.height == 0) and self.scale != 1:
64 | output_width = int(image.width * self.scale)
65 | output_height = int(image.height * self.scale)
66 | elif self.width != 0 and self.height != 0:
67 | output_width = self.width
68 | output_height = self.height
69 |
70 | if image.width != output_width or image.height != output_height:
71 | LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
72 | if self.upscaler == constants.fast_upscaler:
73 | image = image.resize((output_width, output_height), resample=LANCZOS)
74 | else:
75 | image = images.resize_image(0, image, output_width, output_height, self.upscaler)
76 |
77 | if self.prompt == '':
78 | self.prompt = context.get_prompt_by_index()
79 | debug_print('prompt', self.prompt)
80 | elif self.prompt.find('#!org!#') >= 0:
81 | current_prompt = context.get_prompt_by_index()
82 | self.prompt = self.prompt.replace('#!org!#', current_prompt)
83 | debug_print('Prompt', self.prompt)
84 | if self.negative_prompt == '':
85 | self.negative_prompt = context.sdprocessing.negative_prompt
86 | if self.sampler == constants.sampler_default:
87 | self.sampler = context.sdprocessing.sampler_name
88 | if self.scheduler == constants.scheduler_default:
89 | self.scheduler = util.get_scheduler(context.sdprocessing)
90 |
91 | seed, subseed = context.get_seeds()
92 | options = dict(
93 | seed=seed, subseed=subseed,
94 | denoising_strength=self.denoising_strength,
95 | resize_mode=0,
96 | mask=None,
97 | mask_blur=4,
98 | inpainting_fill=1,
99 | inpaint_full_res=True,
100 | inpaint_full_res_padding=32,
101 | inpainting_mask_invert=0,
102 | initial_noise_multiplier=1.0,
103 | prompt=self.prompt,
104 | negative_prompt=self.negative_prompt,
105 | sampler_name=self.sampler,
106 | scheduler=self.scheduler,
107 | batch_size=1,
108 | n_iter=1,
109 | steps=self.steps,
110 | cfg_scale=self.cfg_scale,
111 | width=output_width,
112 | height=output_height,
113 | restore_faces=False,
114 | do_not_save_samples=True,
115 | do_not_save_grid=True,
116 | )
117 | context.add_job()
118 |
119 | if self.checkpoint is not None and self.checkpoint != constants.checkpoint_default:
120 | override_settings = options.get('override_settings', {})
121 | override_settings['sd_model_checkpoint'] = self.checkpoint
122 | options['override_settings'] = override_settings
123 | if self.vae is not None and self.vae != constants.vae_default:
124 | override_settings = options.get('override_settings', {})
125 | override_settings['sd_vae'] = self.vae
126 | options['override_settings'] = override_settings
127 |
128 | if LineartNoise.with_refiner(context):
129 | ln = LineartNoise()
130 | if ln.preprocess(context, None):
131 | controlnet = ln.get_controlnet_args(context)
132 | image = process_img2img_with_controlnet(context, image, options, controlnet=[controlnet])
133 | else:
134 | image = process_img2img(context, image, options=options)
135 | else:
136 | image = process_img2img(context, image, options=options)
137 |
138 | if not self.keep_checkpoint:
139 | debug_print('Rollback model')
140 | context.restore_checkpoint()
141 |
142 | return image
143 |
144 | @staticmethod
145 | def process_callback(self, context, img2img):
146 | ctx = Context.newContext(self, img2img, context.args, 0)
147 | ctx.refiner = self
148 | ln = LineartNoise()
149 | if ln.preprocess(ctx, None):
150 | ln.process(ctx, None)
151 | ln.postprocess(ctx, None)
152 |
153 | def postprocess(self, context: Context, image: Image):
154 | devices.torch_gc()
155 |
--------------------------------------------------------------------------------
/sd_bmab/processors/utils/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.processors.utils.filesave import AfterProcessFileSaver, BeforeProcessFileSaver
2 | from sd_bmab.processors.utils.modelswitch import ApplyModel, RollbackModel
3 | from sd_bmab.processors.utils.checkpoint import CheckPointChanger, CheckPointRestore
4 |
--------------------------------------------------------------------------------
/sd_bmab/processors/utils/checkpoint.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from sd_bmab import constants
4 | from sd_bmab.base.context import Context
5 | from sd_bmab.base.processorbase import ProcessorBase
6 | from sd_bmab.util import debug_print
7 |
8 |
9 | class CheckPointChanger(ProcessorBase):
10 | def __init__(self) -> None:
11 | super().__init__()
12 | self.checkpoint = constants.checkpoint_default
13 | self.vae = constants.vae_default
14 |
15 | def preprocess(self, context: Context, image: Image):
16 | self.checkpoint = context.args['preprocess_checkpoint']
17 | self.vae = context.args['preprocess_vae']
18 | return not (self.checkpoint == constants.checkpoint_default and self.vae == constants.vae_default)
19 |
20 | def process(self, context: Context, image: Image):
21 | debug_print('Change checkpoint', self.checkpoint, self.vae)
22 | context.save_and_apply_checkpoint(self.checkpoint, self.vae)
23 | return image
24 |
25 |
26 | class CheckPointRestore(ProcessorBase):
27 | def __init__(self) -> None:
28 | super().__init__()
29 |
30 | def preprocess(self, context: Context, image: Image):
31 | return True
32 |
33 | def process(self, context: Context, image: Image):
34 | context.restore_checkpoint()
35 | return image
36 |
--------------------------------------------------------------------------------
/sd_bmab/processors/utils/filesave.py:
--------------------------------------------------------------------------------
1 | from PIL import Image
2 |
3 | from modules import shared
4 | from modules import images
5 |
6 | from sd_bmab.base.context import Context
7 | from sd_bmab.base.processorbase import ProcessorBase
8 |
9 |
10 | class BeforeProcessFileSaver(ProcessorBase):
11 | def __init__(self) -> None:
12 | super().__init__()
13 |
14 | def preprocess(self, context: Context, image: Image):
15 | return shared.opts.bmab_save_image_before_process
16 |
17 | def process(self, context: Context, image: Image):
18 | images.save_image(
19 | image, context.sdprocessing.outpath_samples, '',
20 | context.sdprocessing.all_seeds[context.index], context.sdprocessing.all_prompts[context.index],
21 | shared.opts.samples_format, p=context.sdprocessing, suffix='-before-bmab')
22 | return image
23 |
24 | def postprocess(self, context: Context, image: Image):
25 | pass
26 |
27 |
28 | class AfterProcessFileSaver(ProcessorBase):
29 | def __init__(self) -> None:
30 | super().__init__()
31 |
32 | def preprocess(self, context: Context, image: Image):
33 | return shared.opts.bmab_save_image_after_process
34 |
35 | def process(self, context: Context, image: Image):
36 | images.save_image(
37 | image, context.sdprocessing.outpath_samples, '',
38 | context.sdprocessing.all_seeds[context.index], context.sdprocessing.all_prompts[context.index],
39 | shared.opts.samples_format, p=context.sdprocessing, suffix="-after-bmab")
40 | return image
41 |
42 | def postprocess(self, context: Context, image: Image):
43 | pass
44 |
--------------------------------------------------------------------------------
/sd_bmab/processors/utils/modelswitch.py:
--------------------------------------------------------------------------------
1 |
2 | from PIL import Image
3 |
4 | from modules import shared, sd_models
5 |
6 | from sd_bmab.base.context import Context
7 | from sd_bmab.base.processorbase import ProcessorBase
8 |
9 |
10 | base_sd_model = None
11 |
12 |
13 | def change_model(name):
14 | if name is None:
15 | return
16 | info = sd_models.get_closet_checkpoint_match(name)
17 | if info is None:
18 | print(f'Unknown model: {name}')
19 | return
20 | sd_models.reload_model_weights(shared.sd_model, info)
21 |
22 |
23 | class ApplyModel(ProcessorBase):
24 | def __init__(self) -> None:
25 | super().__init__()
26 |
27 | def preprocess(self, context: Context, image: Image):
28 | return shared.opts.bmab_use_specific_model
29 |
30 | def process(self, context: Context, image: Image):
31 | global base_sd_model
32 | base_sd_model = shared.opts.data['sd_model_checkpoint']
33 | change_model(shared.opts.bmab_model)
34 | return image
35 |
36 | def postprocess(self, context: Context, image: Image):
37 | pass
38 |
39 |
40 | class RollbackModel(ProcessorBase):
41 | def __init__(self) -> None:
42 | super().__init__()
43 |
44 | def preprocess(self, context: Context, image: Image):
45 | return shared.opts.bmab_use_specific_model
46 |
47 | def process(self, context: Context, image: Image):
48 | global base_sd_model
49 | if base_sd_model is not None:
50 | change_model(base_sd_model)
51 | base_sd_model = None
52 | return image
53 |
54 | def postprocess(self, context: Context, image: Image):
55 | pass
56 |
--------------------------------------------------------------------------------
/sd_bmab/sd_override/__init__.py:
--------------------------------------------------------------------------------
1 | from sd_bmab.sd_override.samper import override_samplers
2 | from sd_bmab.sd_override.img2img import StableDiffusionProcessingImg2ImgOv
3 | from sd_bmab.sd_override.txt2img import StableDiffusionProcessingTxt2ImgOv
4 |
5 | from modules import processing
6 |
7 | processing.StableDiffusionProcessingImg2Img = StableDiffusionProcessingImg2ImgOv
8 | processing.StableDiffusionProcessingTxt2Img = StableDiffusionProcessingTxt2ImgOv
9 |
10 |
11 | def override_sd_webui():
12 | # override_samplers()
13 | pass
14 |
--------------------------------------------------------------------------------
/sd_bmab/sd_override/img2img.py:
--------------------------------------------------------------------------------
1 | from dataclasses import dataclass
2 |
3 | from modules.processing import StableDiffusionProcessingImg2Img
4 |
5 |
6 | @dataclass(repr=False)
7 | class StableDiffusionProcessingImg2ImgOv(StableDiffusionProcessingImg2Img):
8 | extra_noise: int = 0
9 |
10 | def __post_init__(self):
11 | super().__post_init__()
12 |
13 | def init(self, all_prompts, all_seeds, all_subseeds):
14 | ret = super().init(all_prompts, all_seeds, all_subseeds)
15 | self.extra_generation_params['Hires prompt'] = ''
16 | self.extra_generation_params['Hires negative prompt'] = ''
17 | return ret
18 |
19 | def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts):
20 | return super().sample(conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength, prompts)
21 |
22 |
--------------------------------------------------------------------------------
/sd_bmab/sd_override/samper.py:
--------------------------------------------------------------------------------
1 | import modules
2 | import k_diffusion.sampling
3 | from modules.sd_samplers_kdiffusion import KDiffusionSampler
4 | from modules.sd_samplers import set_samplers
5 | from modules.shared import opts, state
6 | import inspect
7 | from modules import sd_samplers_common
8 | from modules.script_callbacks import ExtraNoiseParams, extra_noise_callback
9 |
10 |
11 | class KDiffusionSamplerBMAB(KDiffusionSampler):
12 |
13 | def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
14 | steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)
15 |
16 | sigmas = self.get_sigmas(p, steps)
17 | sigma_sched = sigmas[steps - t_enc - 1:]
18 |
19 | xi = x + noise * sigma_sched[0]
20 |
21 | if opts.img2img_extra_noise > 0:
22 | p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise
23 | extra_noise_params = ExtraNoiseParams(noise, x, xi)
24 | extra_noise_callback(extra_noise_params)
25 | noise = extra_noise_params.noise
26 | xi += noise * opts.img2img_extra_noise
27 |
28 | if hasattr(p, 'extra_noise') and p.extra_noise > 0:
29 | p.extra_generation_params["Extra noise"] = p.extra_noise
30 | extra_noise_params = ExtraNoiseParams(noise, x, xi)
31 | extra_noise_callback(extra_noise_params)
32 | noise = extra_noise_params.noise
33 | xi += noise * p.extra_noise
34 |
35 | extra_params_kwargs = self.initialize(p)
36 | parameters = inspect.signature(self.func).parameters
37 |
38 | if 'sigma_min' in parameters:
39 | ## last sigma is zero which isn't allowed by DPM Fast & Adaptive so taking value before last
40 | extra_params_kwargs['sigma_min'] = sigma_sched[-2]
41 | if 'sigma_max' in parameters:
42 | extra_params_kwargs['sigma_max'] = sigma_sched[0]
43 | if 'n' in parameters:
44 | extra_params_kwargs['n'] = len(sigma_sched) - 1
45 | if 'sigma_sched' in parameters:
46 | extra_params_kwargs['sigma_sched'] = sigma_sched
47 | if 'sigmas' in parameters:
48 | extra_params_kwargs['sigmas'] = sigma_sched
49 |
50 | if self.config.options.get('brownian_noise', False):
51 | noise_sampler = self.create_noise_sampler(x, sigmas, p)
52 | extra_params_kwargs['noise_sampler'] = noise_sampler
53 |
54 | if self.config.options.get('solver_type', None) == 'heun':
55 | extra_params_kwargs['solver_type'] = 'heun'
56 |
57 | self.model_wrap_cfg.init_latent = x
58 | self.last_latent = x
59 | self.sampler_extra_args = {
60 | 'cond': conditioning,
61 | 'image_cond': image_conditioning,
62 | 'uncond': unconditional_conditioning,
63 | 'cond_scale': p.cfg_scale,
64 | 's_min_uncond': self.s_min_uncond
65 | }
66 |
67 | samples = self.launch_sampling(t_enc + 1, lambda: self.func(self.model_wrap_cfg, xi, extra_args=self.sampler_extra_args, disable=False, callback=self.callback_state, **extra_params_kwargs))
68 |
69 | if self.model_wrap_cfg.padded_cond_uncond:
70 | p.extra_generation_params["Pad conds"] = True
71 |
72 | return samples
73 |
74 |
75 | def override_samplers():
76 | modules.sd_samplers_kdiffusion.samplers_data_k_diffusion = [
77 | modules.sd_samplers_common.SamplerData(label,
78 | lambda model, funcname=funcname: KDiffusionSamplerBMAB(funcname, model),
79 | aliases, options)
80 | for label, funcname, aliases, options in modules.sd_samplers_kdiffusion.samplers_k_diffusion
81 | if callable(funcname) or hasattr(k_diffusion.sampling, funcname)
82 | ]
83 | if hasattr(modules, 'sd_samplers_timesteps'):
84 | modules.sd_samplers.all_samplers = [
85 | *modules.sd_samplers_kdiffusion.samplers_data_k_diffusion,
86 | *modules.sd_samplers_timesteps.samplers_data_timesteps,
87 | ]
88 | else:
89 | modules.sd_samplers.all_samplers = [
90 | *modules.sd_samplers_kdiffusion.samplers_data_k_diffusion,
91 | *modules.sd_samplers_compvis.samplers_data_compvis,
92 | ]
93 | modules.sd_samplers.all_samplers_map = {x.name: x for x in modules.sd_samplers.all_samplers}
94 |
--------------------------------------------------------------------------------
/sd_bmab/sd_override/sd_models.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | from modules import sd_models
4 | from modules import shared
5 |
6 |
7 | def bmab_list_models():
8 | sd_models.checkpoints_list.clear()
9 | sd_models.checkpoint_aliases.clear()
10 |
11 | cmd_ckpt = shared.cmd_opts.ckpt
12 | if shared.cmd_opts.no_download_sd_model or cmd_ckpt != shared.sd_model_file or os.path.exists(cmd_ckpt):
13 | model_url = None
14 | else:
15 | model_url = f"{shared.hf_endpoint}/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
16 |
17 | model_list = sd_models.modelloader.load_models(model_path=sd_models.model_path, model_url=model_url, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name="v1-5-pruned-emaonly.safetensors", ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
18 |
19 | second_path = shared.opts.data.get('bmab_additional_checkpoint_path', '')
20 | print(f'second path {second_path}')
21 | if os.path.exists(second_path):
22 | print(f'load checkpoint from {second_path}')
23 | model_list_seconds = sd_models.modelloader.load_models(model_path=second_path, model_url=model_url, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"], download_name="v1-5-pruned-emaonly.safetensors", ext_blacklist=[".vae.ckpt", ".vae.safetensors"])
24 | length = len(sd_models.model_path)
25 | temp = [(x[length:], x) for x in model_list]
26 | length = len(second_path)
27 | temp.extend([(x[length:], x) for x in model_list_seconds])
28 | model_list = [x[1] for x in sorted(temp, key=lambda x: x[0])]
29 |
30 | if os.path.exists(cmd_ckpt):
31 | checkpoint_info = sd_models.CheckpointInfo(cmd_ckpt)
32 | checkpoint_info.register()
33 | shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title
34 | elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
35 | print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {sd_models.model_path}: {cmd_ckpt}", file=sys.stderr)
36 |
37 | for filename in model_list:
38 | checkpoint_info = sd_models.CheckpointInfo(filename)
39 | checkpoint_info.register()
40 |
41 |
42 | def override():
43 | bmab_list_models()
44 | sd_models.list_models = bmab_list_models
45 |
46 |
--------------------------------------------------------------------------------
/sd_bmab/util/installhelper.py:
--------------------------------------------------------------------------------
1 | import shutil
2 |
3 | import launch
4 | from modules import launch_utils
5 | import torch
6 | import platform
7 |
8 |
9 | groundingdino_install_info = {
10 | '2.1.2+cu121:3.10:linux:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp310-cp310-manylinux_2_34_x86_64.whl',
11 | '2.1.2+cu121:3.10:windows:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp310-cp310-win_amd64.whl',
12 | '2.1.2+cu121:3.11:linux:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp311-cp311-manylinux_2_34_x86_64.whl',
13 | '2.1.2+cu121:3.11:windows:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp311-cp311-win_amd64.whl',
14 | '2.1.2+cu121:3.9:linux:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp39-cp39-manylinux_2_34_x86_64.whl',
15 | '2.1.2+cu121:3.9:windows:amd64': 'https://github.com/Bing-su/GroundingDINO/releases/download/v23.9.27/groundingdino-23.9.27+torch2.1.0.cu121-cp39-cp39-win_amd64.whl',
16 | }
17 |
18 | groundingdino_install_replacement = {
19 | '2.1.0+cu121:3.10:linux:amd64': '2.1.2+cu121:3.10:linux:amd64',
20 | '2.1.1+cu121:3.10:linux:amd64': '2.1.2+cu121:3.10:linux:amd64',
21 | }
22 |
23 |
24 | available_packages = ['GroundingDINO']
25 |
26 |
27 | def install_groudingdino(url):
28 | launch_utils.run(f'{launch_utils.python} -m pip uninstall --yes groundingdino', live=True)
29 | launch_utils.run(f'{launch_utils.python} -m pip uninstall --yes pycocotools', live=True)
30 |
31 | print('Install pycocotools')
32 | launch.run_pip('install pycocotools', 'sd-webui-bmab requirement: pycocotools')
33 |
34 | print(f'Install groundingdino from {url}. It takes minutes.')
35 | launch.run_pip(f'install {url}', 'sd-webui-bmab requirement: groundingdino')
36 | print('Done.')
37 |
38 |
39 | def get_condition():
40 | torch_version = torch.__version__
41 | pv = platform.python_version_tuple()
42 | system = 'windows' if platform.system() == 'Windows' else 'linux'
43 | machine = 'amd64' if platform.machine() == 'AMD64' else 'x86_64'
44 | return f'{torch_version}:{pv[0]}.{pv[1]}:{system}:{machine}'
45 |
46 |
47 | def get_temporary():
48 | if platform.system() == 'Windows':
49 | return 'c:\\temp'
50 | return '/temp'
51 |
52 |
53 | def install(pkg_name, dd_pkg, markdown_install):
54 | groundingdino_cuda_name = 'GroundingDINO for CUDA'
55 | groundingdino_selected = pkg_name
56 |
57 | def add_new_available(cond):
58 | msg = f'GroundingDINO for CUDA Found {cond}. Please select GroudingDINO or {groundingdino_cuda_name}'
59 | newname = f'{groundingdino_cuda_name}-{cond}'
60 | if newname not in available_packages:
61 | available_packages.append(newname)
62 | return msg, newname
63 |
64 | def install_normal_groundingdino(c):
65 | url = 'https://github.com/IDEA-Research/GroundingDINO'
66 | launch_utils.run(f'{launch_utils.python} -m pip uninstall --yes groundingdino', live=True)
67 | launch_utils.run(f'{launch_utils.python} -m pip uninstall --yes pycocotools', live=True)
68 | launch_utils.run(f'{launch_utils.python} -m pip install pycocotools', live=True)
69 | if platform.system() == 'Windows':
70 | launch_utils.run(f'{launch_utils.git} clone {url} c:\\Temp\\groundingdino', live=True)
71 | launch_utils.run(f'cd c:\\Temp\\groundingdino && {launch_utils.python} -m pip install -e .', live=True)
72 | shutil.rmtree('c:\\Temp\\groundingdino', ignore_errors=True)
73 | else:
74 | launch_utils.run(f'{launch_utils.git} clone {url} /temp/groundingdino', live=True)
75 | launch_utils.run(f'cd /temp/groundingdino && {launch_utils.python} -m pip install -e .', live=True)
76 | shutil.rmtree('rm -rf /temp/groundingdino', ignore_errors=True)
77 |
78 | return f'Nothing found for this cuda system {c}. Software version of GroundingDINO installed (VERY SLOW!!!)'
79 |
80 | def cuda_in_available_packages():
81 | for x in available_packages:
82 | if x.startswith(groundingdino_cuda_name):
83 | return True
84 | return False
85 |
86 | if pkg_name == 'GroundingDINO':
87 | cond = get_condition()
88 | if cuda_in_available_packages():
89 | msg = install_normal_groundingdino(cond)
90 | else:
91 | replacement = groundingdino_install_replacement.get(cond)
92 | if replacement is not None:
93 | msg, groundingdino_selected = add_new_available(cond)
94 | else:
95 | groudingdino_for_cuda = groundingdino_install_info.get(cond)
96 | if groudingdino_for_cuda is None:
97 | msg = install_normal_groundingdino(cond)
98 | else:
99 | msg, groundingdino_selected = add_new_available(cond)
100 | elif pkg_name.startswith(groundingdino_cuda_name):
101 | groudingdino_for_cuda_cond = pkg_name[len(groundingdino_cuda_name)+1:]
102 | groudingdino_for_cuda = groundingdino_install_info.get(groudingdino_for_cuda_cond)
103 | if groudingdino_for_cuda is not None:
104 | install_groudingdino(groudingdino_for_cuda)
105 | msg = f'{groundingdino_cuda_name} installed. {groudingdino_for_cuda}'
106 | groundingdino_selected = f'{groundingdino_cuda_name}-{groudingdino_for_cuda_cond}'
107 | else:
108 | msg = 'Nothing installed.'
109 | else:
110 | msg = 'Nothing installed.'
111 |
112 | return {
113 | dd_pkg: {
114 | 'choices': available_packages,
115 | 'value': groundingdino_selected,
116 | '__type__': 'update'
117 | },
118 | markdown_install: {
119 | 'value': msg,
120 | '__type__': 'update'
121 | }
122 | }
123 |
124 |
--------------------------------------------------------------------------------