forked from AIrjen/OneButtonPrompt
-
Notifications
You must be signed in to change notification settings - Fork 0
/
call_img2img.py
219 lines (163 loc) · 8.73 KB
/
call_img2img.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
import os
import requests
import io
import base64
import uuid
from PIL import Image, PngImagePlugin
from modules import shared
from model_lists import *
import time
def call_img2img(imagelocation,originalimage, originalpnginfo ="", apiurl="http://127.0.0.1:7860",filename="", prompt = "", negativeprompt = "", img2imgsamplingsteps = "20", img2imgcfg = "7", img2imgsamplingmethod = "DPM++ SDE Karras", img2imgupscaler = "R-ESRGAN 4x+", img2imgmodel = "currently selected model", denoising_strength = "0.3", scale = "2", padding = "64",upscalescript="SD upscale",usdutilewidth = "512", usdutileheight = "0", usdumaskblur = "8", usduredraw ="Linear", usduSeamsfix = "None", usdusdenoise = "0.35", usduswidth = "64", usduspadding ="32", usdusmaskblur = "8",controlnetenabled=False, controlnetmodel="",controlnetblockymode=False):
negativepromptfound = 0
#params to stay the same
url = apiurl
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
outputimg2imgfolder = os.path.join(script_dir, "./automated_outputs/img2img/" )
outputimg2imgfolder.replace("./", "/")
if(filename==""):
filename = str(uuid.uuid4())
outputimg2imgpng = '.png'
outputimg2imgFull = '{}{}{}'.format(outputimg2imgfolder,filename,outputimg2imgpng)
encodedstringlist = []
# need to convert the values to the correct index number for Ultimate SD Upscaler
redrawmodelist =["Linear","Chess","None"]
seamsfixmodelist = ["None","Band pass","Half tile offset pass","Half tile offset pass + intersections"]
usduredrawint = int(redrawmodelist.index(usduredraw))
seamsfixmodeint = int(seamsfixmodelist.index(usduSeamsfix))
#rest of prompt things
sampler_index = img2imgsamplingmethod
steps = img2imgsamplingsteps
cfg_scale = img2imgcfg
with open(imagelocation, "rb") as image_file:
encoded_string = base64.b64encode(image_file.read())
encodedstringlist.append(encoded_string.decode('utf-8'))
# If we don't have a prompt, get it from the original image file
# This is used when only_upscale is activated
if(prompt==""):
with open(originalimage, "rb") as originalimage_file:
originalencoded_string = base64.b64encode(originalimage_file.read())
encodedstring2 = originalencoded_string.decode('utf-8')
# get prompt from picture
png_payload = {
"image": encodedstring2
}
response3 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
pnginfo = str(response3.json().get("info"))
prompt = pnginfo[:pnginfo.rfind("Steps")]
if(prompt.rfind("Negative prompt") != -1):
prompt = prompt[:prompt.rfind("Negative prompt")]
negativepromptfound = 1
if(negativepromptfound == 1):
negativeprompt = pnginfo[:pnginfo.rfind("Steps")]
negativeprompt = negativeprompt.replace(prompt,"")
# set the automatic upscale
checkprompt = prompt.lower()
if(img2imgupscaler != "automatic"):
upscaler = img2imgupscaler
else:
upscalerlist = get_upscalers_for_img2img()
# on automatic, make some choices about what upscaler to use
# photos, prefer 4x ultrasharp
# anime, cartoon or drawing, go for R-ESRGAN 4x+ Anime6B
# else, R-ESRGAN 4x+"
if("hoto" in checkprompt and "4x-UltraSharp" in upscalerlist):
upscaler = "4x-UltraSharp"
elif("anime" in checkprompt or "cartoon" in checkprompt or "draw" in checkprompt or "vector" in checkprompt or "cel shad" in checkprompt or "visual novel" in checkprompt):
upscaler = "R-ESRGAN 4x+ Anime6B"
else:
upscaler = "R-ESRGAN 4x+"
if(upscaler== "4x-UltraSharp"):
denoising_strength = "0.35"
if(upscaler== "R-ESRGAN 4x+ Anime6B"):
denoising_strength = "0.6" # 0.6 is fine for the anime upscaler
if(upscaler== "R-ESRGAN 4x+"):
denoising_strength = "0.5" # default 0.6 is a lot and changes a lot of details
#wierd blocky mode comes up when the treshold is set way too high and the denoising strenght is strong
if(controlnetblockymode==True):
treshold = int(padding)
if(float(denoising_strength) < 0.65):
denoising_strength = "0.65"
else:
treshold = 1
payload = {
"resize_mode": 0,
"denoising_strength": denoising_strength,
"sampler_index": sampler_index,
"batch_size": "1",
"n_iter": "1",
"prompt": prompt,
"negative_prompt": negativeprompt,
"steps": steps,
"cfg_scale": cfg_scale,
#"width": width,
#"height": height,
"include_init_images": "true",
"init_images": encodedstringlist,
}
if(img2imgmodel != "currently selected model"):
payload.update({"sd_model": img2imgmodel})
#
# https://github.com/Mikubill/sd-webui-controlnet/wiki/API
#
if(controlnetenabled==True and controlnetmodel!=""):
payload.update({"alwayson_scripts": {
"controlnet": {
"args": [
{
"module": "tile_resample",
"model": controlnetmodel, # control_v11f1e_sd15_tile [a371b31b]
#"input_image": encodedstringlist,
"control_mode": 2, #"ControlNet is more important" : the controlnet model has more impact than the prompt
#"resize_mode": 0
"threshold_a": treshold
}
]
}
}
})
if(upscalescript=="SD upscale"):
payload.update({"script_name": upscalescript})
payload.update({"script_args": ["",int(padding),upscaler,round(float(scale),1)]})
if(upscalescript=="Ultimate SD upscale"):
upscaler_index = [x.name.lower() for x in shared.sd_upscalers].index(upscaler.lower())
payload.update({"script_name": upscalescript})
payload.update({"script_args": ["",int(usdutilewidth),int(usdutileheight),int(usdumaskblur),int(padding), int(usduswidth), round(float(usdusdenoise),2),int(usduspadding),
upscaler_index,True,usduredrawint,False,int(usdusmaskblur),
seamsfixmodeint,2,"","",round(float(scale),1)]})
# Ultimate SD Upscale params:
#_, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
# upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
# seams_fix_type, target_size_type, custom_width, custom_height, custom_scale):
# target_size_type = 2
# custom_scale = 2
r = []
# If we don't get an image back, we want to retry a few times. Max 3 times
for i in range(4):
response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
r = response.json()
if('images' in r):
break # this means if we have the images object, then we "break" out of the for loop.
else:
if(i == 3):
print("If this keeps happening: Is WebUI started with --api enabled?")
print("")
raise ValueError("API has not been responding after several retries. Stopped processing.")
print("")
print("We haven't received an image from the API. Maybe something went wrong. Will retry after waiting a bit.")
time.sleep(10 * (i+1) ) # incremental waiting time
for i in r['images']:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
if(originalpnginfo==""):
png_payload = {
"image": "data:image/png;base64," + i
}
#print("and here!")
#print(png_payload)
response2 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
#print("here!")
#print(response2)
pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.json().get("info"))
originalpnginfo = pnginfo
image.save(outputimg2imgFull, pnginfo=originalpnginfo)
return [outputimg2imgFull,originalpnginfo]