forked from philz1337x/clarity-upscaler
-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict.py
181 lines (161 loc) · 6.93 KB
/
predict.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
from modules import timer
from modules import launch_utils
from modules import initialize_util
from modules import initialize
from fastapi import FastAPI
import base64
import os, sys, json
from PIL import Image
import uuid
from io import BytesIO
from cog import BasePredictor, Input, Path
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
os.environ['IGNORE_CMD_ARGS_ERRORS'] = '1'
startup_timer = timer.startup_timer
startup_timer.record("launcher")
initialize.imports()
initialize.check_versions()
initialize.initialize()
app = FastAPI()
initialize_util.setup_middleware(app)
from modules.api.api import Api
from modules.call_queue import queue_lock
self.api = Api(app, queue_lock)
from modules import script_callbacks
script_callbacks.before_ui_callback()
script_callbacks.app_started_callback(None, app)
print(f"Startup time: {startup_timer.summary()}.")
def predict(
self,
image: Path = Input(description="input image"),
prompt: str = Input(description="Prompt", default="masterpiece, best quality, highres, <lora:more_details:0.5> <lora:SDXLrender_v2.0:1>"),
negative_prompt: str = Input(description="Negative Prompt", default="(worst quality, low quality, normal quality:2) JuggernautNegative-neg"),
scale_factor: float = Input(
description="Scale factor", default=2
),
dynamic: float = Input(
description="HDR, try from 3 - 9", ge=1, le=50, default=6
),
creativity: float = Input(
description="Creativity, try from 0.3 - 0.9", ge=0, le=1, default=0.35
),
resemblance: float = Input(
description="Resemblance, try from 0.3 - 1.6", default=0.6
),
tiling_width: int = Input(
description="Fractality, set lower tile width for a high Fractality",
choices=[16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256],
default=112
),
tiling_height: int = Input(
description="Fractality, set lower tile height for a high Fractality",
choices=[16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256],
default=144
),
sd_model: str = Input(
description="Stable Diffusion model checkpoint",
choices=['epicrealism_naturalSinRC1VAE.safetensors [84d76a0328]', 'juggernaut_reborn.safetensors [338b85bc4f]', 'juggernaut_final.safetensors'],
default="juggernaut_reborn.safetensors [338b85bc4f]",
),
scheduler: str = Input(
description="scheduler",
choices=['DPM++ 2M Karras', 'DPM++ SDE Karras', 'DPM++ 2M SDE Exponential', 'DPM++ 2M SDE Karras', 'Euler a', 'Euler', 'LMS', 'Heun', 'DPM2', 'DPM2 a', 'DPM++ 2S a', 'DPM++ 2M', 'DPM++ SDE', 'DPM++ 2M SDE', 'DPM++ 2M SDE Heun', 'DPM++ 2M SDE Heun Karras', 'DPM++ 2M SDE Heun Exponential', 'DPM++ 3M SDE', 'DPM++ 3M SDE Karras', 'DPM++ 3M SDE Exponential', 'DPM fast', 'DPM adaptive', 'LMS Karras', 'DPM2 Karras', 'DPM2 a Karras', 'DPM++ 2S a Karras', 'Restart', 'DDIM', 'PLMS', 'UniPC'],
default="DPM++ 3M SDE Karras",
),
num_inference_steps: int = Input(
description="Number of denoising steps", ge=1, le=100, default=18
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed", default=1337
)
) -> list[Path]:
"""Run a single prediction on the model"""
from modules.api.models import StableDiffusionImg2ImgProcessingAPI
with open(image, "rb") as image_file:
img_base64 = base64.b64encode(image_file.read()).decode()
payload = {
"override_settings": {
"sd_model_checkpoint": sd_model,
"sd_vae": "vae-ft-mse-840000-ema-pruned.safetensors",
"CLIP_stop_at_last_layers": 1,
},
"override_settings_restore_afterwards": False,
"init_images": [img_base64],
"prompt": prompt,
"negative_prompt": negative_prompt,
"steps": num_inference_steps,
"cfg_scale": dynamic,
"seed": seed,
"do_not_save_samples": True,
"sampler_name": scheduler,
"denoising_strength": creativity,
"alwayson_scripts": {
"Tiled Diffusion": {
"args": [
True,
"MultiDiffusion",
True,
True,
1,
1,
tiling_width,
tiling_height,
4,
8,
"4x-UltraSharp",
scale_factor,
False,
0,
0.0,
3,
]
},
"Tiled VAE": {
"args": [
True,
3072,
192,
True,
True,
True,
True,
]
},
"controlnet": {
"args": [
{
"enabled": True,
"module": "tile_resample",
"model": "control_v11f1e_sd15_tile",
"weight": resemblance,
"image": img_base64,
"resize_mode": 1,
"lowvram": False,
"downsample": 1.0,
"guidance_start": 0.0,
"guidance_end": 1.0,
"control_mode": 1,
"pixel_perfect": True,
"threshold_a": 1,
"threshold_b": 1,
"save_detected_map": False,
"processor_res": 512,
}
]
}
}
}
req = StableDiffusionImg2ImgProcessingAPI(**payload)
resp = self.api.img2imgapi(req)
info = json.loads(resp.info)
outputs = []
for i, image in enumerate(resp.images):
seed = info.get("all_seeds", [])[i] or "unknown_seed"
gen_bytes = BytesIO(base64.b64decode(image))
filename = f"{seed}-{uuid.uuid1()}.png"
with open(filename, "wb") as f:
f.write(gen_bytes.getvalue())
outputs.append(Path(filename))
return outputs