File size: 8,938 Bytes
f4a41d8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
220
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]