Upload dc.py
Browse files
dc.py
CHANGED
@@ -59,202 +59,60 @@ from stablepy import logger
|
|
59 |
logger.setLevel(logging.DEBUG)
|
60 |
|
61 |
from env import (
|
62 |
-
HF_TOKEN,
|
63 |
CIVITAI_API_KEY, HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
|
64 |
HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
|
65 |
HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
PREPROCESSOR_CONTROLNET = {
|
72 |
-
"openpose": [
|
73 |
-
"Openpose",
|
74 |
-
"None",
|
75 |
-
],
|
76 |
-
"scribble": [
|
77 |
-
"HED",
|
78 |
-
"PidiNet",
|
79 |
-
"None",
|
80 |
-
],
|
81 |
-
"softedge": [
|
82 |
-
"PidiNet",
|
83 |
-
"HED",
|
84 |
-
"HED safe",
|
85 |
-
"PidiNet safe",
|
86 |
-
"None",
|
87 |
-
],
|
88 |
-
"segmentation": [
|
89 |
-
"UPerNet",
|
90 |
-
"None",
|
91 |
-
],
|
92 |
-
"depth": [
|
93 |
-
"DPT",
|
94 |
-
"Midas",
|
95 |
-
"None",
|
96 |
-
],
|
97 |
-
"normalbae": [
|
98 |
-
"NormalBae",
|
99 |
-
"None",
|
100 |
-
],
|
101 |
-
"lineart": [
|
102 |
-
"Lineart",
|
103 |
-
"Lineart coarse",
|
104 |
-
"Lineart (anime)",
|
105 |
-
"None",
|
106 |
-
"None (anime)",
|
107 |
-
],
|
108 |
-
"lineart_anime": [
|
109 |
-
"Lineart",
|
110 |
-
"Lineart coarse",
|
111 |
-
"Lineart (anime)",
|
112 |
-
"None",
|
113 |
-
"None (anime)",
|
114 |
-
],
|
115 |
-
"shuffle": [
|
116 |
-
"ContentShuffle",
|
117 |
-
"None",
|
118 |
-
],
|
119 |
-
"canny": [
|
120 |
-
"Canny",
|
121 |
-
"None",
|
122 |
-
],
|
123 |
-
"mlsd": [
|
124 |
-
"MLSD",
|
125 |
-
"None",
|
126 |
-
],
|
127 |
-
"ip2p": [
|
128 |
-
"ip2p"
|
129 |
-
],
|
130 |
-
"recolor": [
|
131 |
-
"Recolor luminance",
|
132 |
-
"Recolor intensity",
|
133 |
-
"None",
|
134 |
-
],
|
135 |
-
"tile": [
|
136 |
-
"Mild Blur",
|
137 |
-
"Moderate Blur",
|
138 |
-
"Heavy Blur",
|
139 |
-
"None",
|
140 |
-
],
|
141 |
-
}
|
142 |
-
|
143 |
-
TASK_STABLEPY = {
|
144 |
-
'txt2img': 'txt2img',
|
145 |
-
'img2img': 'img2img',
|
146 |
-
'inpaint': 'inpaint',
|
147 |
-
# 'canny T2I Adapter': 'sdxl_canny_t2i', # NO HAVE STEP CALLBACK PARAMETERS SO NOT WORKS WITH DIFFUSERS 0.29.0
|
148 |
-
# 'sketch T2I Adapter': 'sdxl_sketch_t2i',
|
149 |
-
# 'lineart T2I Adapter': 'sdxl_lineart_t2i',
|
150 |
-
# 'depth-midas T2I Adapter': 'sdxl_depth-midas_t2i',
|
151 |
-
# 'openpose T2I Adapter': 'sdxl_openpose_t2i',
|
152 |
-
'openpose ControlNet': 'openpose',
|
153 |
-
'canny ControlNet': 'canny',
|
154 |
-
'mlsd ControlNet': 'mlsd',
|
155 |
-
'scribble ControlNet': 'scribble',
|
156 |
-
'softedge ControlNet': 'softedge',
|
157 |
-
'segmentation ControlNet': 'segmentation',
|
158 |
-
'depth ControlNet': 'depth',
|
159 |
-
'normalbae ControlNet': 'normalbae',
|
160 |
-
'lineart ControlNet': 'lineart',
|
161 |
-
'lineart_anime ControlNet': 'lineart_anime',
|
162 |
-
'shuffle ControlNet': 'shuffle',
|
163 |
-
'ip2p ControlNet': 'ip2p',
|
164 |
-
'optical pattern ControlNet': 'pattern',
|
165 |
-
'recolor ControlNet': 'recolor',
|
166 |
-
'tile ControlNet': 'tile',
|
167 |
-
}
|
168 |
-
|
169 |
-
TASK_MODEL_LIST = list(TASK_STABLEPY.keys())
|
170 |
-
|
171 |
-
UPSCALER_DICT_GUI = {
|
172 |
-
None: None,
|
173 |
-
"Lanczos": "Lanczos",
|
174 |
-
"Nearest": "Nearest",
|
175 |
-
'Latent': 'Latent',
|
176 |
-
'Latent (antialiased)': 'Latent (antialiased)',
|
177 |
-
'Latent (bicubic)': 'Latent (bicubic)',
|
178 |
-
'Latent (bicubic antialiased)': 'Latent (bicubic antialiased)',
|
179 |
-
'Latent (nearest)': 'Latent (nearest)',
|
180 |
-
'Latent (nearest-exact)': 'Latent (nearest-exact)',
|
181 |
-
"RealESRGAN_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
|
182 |
-
"RealESRNet_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
|
183 |
-
"RealESRGAN_x4plus_anime_6B": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
|
184 |
-
"RealESRGAN_x2plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
|
185 |
-
"realesr-animevideov3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
|
186 |
-
"realesr-general-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
|
187 |
-
"realesr-general-wdn-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
|
188 |
-
"4x-UltraSharp": "https://huggingface.co/Shandypur/ESRGAN-4x-UltraSharp/resolve/main/4x-UltraSharp.pth",
|
189 |
-
"4x_foolhardy_Remacri": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
|
190 |
-
"Remacri4xExtraSmoother": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/Remacri%204x%20ExtraSmoother.pth",
|
191 |
-
"AnimeSharp4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/AnimeSharp%204x.pth",
|
192 |
-
"lollypop": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/lollypop.pth",
|
193 |
-
"RealisticRescaler4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/RealisticRescaler%204x.pth",
|
194 |
-
"NickelbackFS4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/NickelbackFS%204x.pth"
|
195 |
-
}
|
196 |
-
|
197 |
-
UPSCALER_KEYS = list(UPSCALER_DICT_GUI.keys())
|
198 |
-
|
199 |
-
|
200 |
-
def get_model_list(directory_path):
|
201 |
-
model_list = []
|
202 |
-
valid_extensions = {'.ckpt', '.pt', '.pth', '.safetensors', '.bin'}
|
203 |
-
|
204 |
-
for filename in os.listdir(directory_path):
|
205 |
-
if os.path.splitext(filename)[1] in valid_extensions:
|
206 |
-
# name_without_extension = os.path.splitext(filename)[0]
|
207 |
-
file_path = os.path.join(directory_path, filename)
|
208 |
-
# model_list.append((name_without_extension, file_path))
|
209 |
-
model_list.append(file_path)
|
210 |
-
print('\033[34mFILE: ' + file_path + '\033[0m')
|
211 |
-
return model_list
|
212 |
|
213 |
## BEGIN MOD
|
214 |
from modutils import (to_list, list_uniq, list_sub, get_model_id_list, get_tupled_embed_list,
|
215 |
get_tupled_model_list, get_lora_model_list, download_private_repo, download_things)
|
216 |
|
217 |
# - **Download Models**
|
218 |
-
download_model = ", ".join(
|
219 |
# - **Download VAEs**
|
220 |
-
download_vae = ", ".join(
|
221 |
# - **Download LoRAs**
|
222 |
-
download_lora = ", ".join(
|
223 |
|
224 |
-
#download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO,
|
225 |
-
download_private_repo(HF_VAE_PRIVATE_REPO,
|
226 |
|
227 |
-
load_diffusers_format_model = list_uniq(
|
228 |
## END MOD
|
229 |
|
230 |
# Download stuffs
|
231 |
for url in [url.strip() for url in download_model.split(',')]:
|
232 |
if not os.path.exists(f"./models/{url.split('/')[-1]}"):
|
233 |
-
download_things(
|
234 |
for url in [url.strip() for url in download_vae.split(',')]:
|
235 |
if not os.path.exists(f"./vaes/{url.split('/')[-1]}"):
|
236 |
-
download_things(
|
237 |
for url in [url.strip() for url in download_lora.split(',')]:
|
238 |
if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
|
239 |
-
download_things(
|
240 |
|
241 |
# Download Embeddings
|
242 |
for url_embed in download_embeds:
|
243 |
if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
|
244 |
-
download_things(
|
245 |
|
246 |
# Build list models
|
247 |
-
embed_list = get_model_list(
|
248 |
-
model_list = get_model_list(
|
249 |
model_list = load_diffusers_format_model + model_list
|
250 |
## BEGIN MOD
|
251 |
lora_model_list = get_lora_model_list()
|
252 |
-
vae_model_list = get_model_list(
|
253 |
vae_model_list.insert(0, "None")
|
254 |
|
255 |
-
#download_private_repo(HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO,
|
256 |
-
#download_private_repo(HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
|
257 |
-
embed_sdxl_list = get_model_list(
|
258 |
|
259 |
def get_embed_list(pipeline_name):
|
260 |
return get_tupled_embed_list(embed_sdxl_list if pipeline_name == "StableDiffusionXLPipeline" else embed_list)
|
@@ -262,99 +120,13 @@ def get_embed_list(pipeline_name):
|
|
262 |
|
263 |
print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
|
264 |
|
265 |
-
msg_inc_vae = (
|
266 |
-
"Use the right VAE for your model to maintain image quality. The wrong"
|
267 |
-
" VAE can lead to poor results, like blurriness in the generated images."
|
268 |
-
)
|
269 |
-
|
270 |
-
SDXL_TASK = [k for k, v in TASK_STABLEPY.items() if v in SDXL_TASKS]
|
271 |
-
SD_TASK = [k for k, v in TASK_STABLEPY.items() if v in SD15_TASKS]
|
272 |
-
FLUX_TASK = list(TASK_STABLEPY.keys())[:3] + [k for k, v in TASK_STABLEPY.items() if v in FLUX_CN_UNION_MODES.keys()]
|
273 |
-
|
274 |
-
MODEL_TYPE_TASK = {
|
275 |
-
"SD 1.5": SD_TASK,
|
276 |
-
"SDXL": SDXL_TASK,
|
277 |
-
"FLUX": FLUX_TASK,
|
278 |
-
}
|
279 |
-
|
280 |
-
MODEL_TYPE_CLASS = {
|
281 |
-
"diffusers:StableDiffusionPipeline": "SD 1.5",
|
282 |
-
"diffusers:StableDiffusionXLPipeline": "SDXL",
|
283 |
-
"diffusers:FluxPipeline": "FLUX",
|
284 |
-
}
|
285 |
-
|
286 |
-
POST_PROCESSING_SAMPLER = ["Use same sampler"] + scheduler_names[:-2]
|
287 |
-
|
288 |
-
def extract_parameters(input_string):
|
289 |
-
parameters = {}
|
290 |
-
input_string = input_string.replace("\n", "")
|
291 |
-
|
292 |
-
if "Negative prompt:" not in input_string:
|
293 |
-
if "Steps:" in input_string:
|
294 |
-
input_string = input_string.replace("Steps:", "Negative prompt: Steps:")
|
295 |
-
else:
|
296 |
-
print("Invalid metadata")
|
297 |
-
parameters["prompt"] = input_string
|
298 |
-
return parameters
|
299 |
-
|
300 |
-
parm = input_string.split("Negative prompt:")
|
301 |
-
parameters["prompt"] = parm[0].strip()
|
302 |
-
if "Steps:" not in parm[1]:
|
303 |
-
print("Steps not detected")
|
304 |
-
parameters["neg_prompt"] = parm[1].strip()
|
305 |
-
return parameters
|
306 |
-
parm = parm[1].split("Steps:")
|
307 |
-
parameters["neg_prompt"] = parm[0].strip()
|
308 |
-
input_string = "Steps:" + parm[1]
|
309 |
-
|
310 |
-
# Extracting Steps
|
311 |
-
steps_match = re.search(r'Steps: (\d+)', input_string)
|
312 |
-
if steps_match:
|
313 |
-
parameters['Steps'] = int(steps_match.group(1))
|
314 |
-
|
315 |
-
# Extracting Size
|
316 |
-
size_match = re.search(r'Size: (\d+x\d+)', input_string)
|
317 |
-
if size_match:
|
318 |
-
parameters['Size'] = size_match.group(1)
|
319 |
-
width, height = map(int, parameters['Size'].split('x'))
|
320 |
-
parameters['width'] = width
|
321 |
-
parameters['height'] = height
|
322 |
-
|
323 |
-
# Extracting other parameters
|
324 |
-
other_parameters = re.findall(r'(\w+): (.*?)(?=, \w+|$)', input_string)
|
325 |
-
for param in other_parameters:
|
326 |
-
parameters[param[0]] = param[1].strip('"')
|
327 |
-
|
328 |
-
return parameters
|
329 |
-
|
330 |
-
def get_model_type(repo_id: str):
|
331 |
-
api = HfApi(token=os.environ.get("HF_TOKEN")) # if use private or gated model
|
332 |
-
default = "SD 1.5"
|
333 |
-
try:
|
334 |
-
model = api.model_info(repo_id=repo_id, timeout=5.0)
|
335 |
-
tags = model.tags
|
336 |
-
for tag in tags:
|
337 |
-
if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
|
338 |
-
except Exception:
|
339 |
-
return default
|
340 |
-
return default
|
341 |
-
|
342 |
## BEGIN MOD
|
343 |
class GuiSD:
|
344 |
-
def __init__(self):
|
345 |
self.model = None
|
346 |
-
|
347 |
-
|
348 |
-
self.
|
349 |
-
base_model_id="Lykon/dreamshaper-8",
|
350 |
-
task_name="txt2img",
|
351 |
-
vae_model=None,
|
352 |
-
type_model_precision=torch.float16,
|
353 |
-
retain_task_model_in_cache=False,
|
354 |
-
device="cpu",
|
355 |
-
)
|
356 |
-
self.model.load_beta_styles()
|
357 |
-
#self.model.device = torch.device("cpu") #
|
358 |
|
359 |
def infer_short(self, model, pipe_params, progress=gr.Progress(track_tqdm=True)):
|
360 |
#progress(0, desc="Start inference...")
|
@@ -368,31 +140,86 @@ class GuiSD:
|
|
368 |
return img
|
369 |
|
370 |
def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
|
371 |
-
|
372 |
-
#yield f"Loading model: {model_name}"
|
373 |
-
|
374 |
vae_model = vae_model if vae_model != "None" else None
|
375 |
model_type = get_model_type(model_name)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
376 |
|
377 |
if vae_model:
|
378 |
vae_type = "SDXL" if "sdxl" in vae_model.lower() else "SD 1.5"
|
379 |
if model_type != vae_type:
|
380 |
-
gr.Warning(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
381 |
|
382 |
-
self.model.device = torch.device("cpu")
|
383 |
-
dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
|
384 |
-
|
385 |
-
self.model.load_pipe(
|
386 |
-
model_name,
|
387 |
-
task_name=TASK_STABLEPY[task],
|
388 |
-
vae_model=vae_model if vae_model != "None" else None,
|
389 |
-
type_model_precision=dtype_model,
|
390 |
-
retain_task_model_in_cache=False,
|
391 |
-
)
|
392 |
#yield f"Model loaded: {model_name}"
|
393 |
|
394 |
#@spaces.GPU
|
395 |
-
|
396 |
def generate_pipeline(
|
397 |
self,
|
398 |
prompt,
|
@@ -497,23 +324,24 @@ class GuiSD:
|
|
497 |
mode_ip2,
|
498 |
scale_ip2,
|
499 |
pag_scale,
|
500 |
-
#progress=gr.Progress(track_tqdm=True),
|
501 |
):
|
502 |
-
|
503 |
-
|
|
|
504 |
vae_model = vae_model if vae_model != "None" else None
|
505 |
loras_list = [lora1, lora2, lora3, lora4, lora5]
|
506 |
vae_msg = f"VAE: {vae_model}" if vae_model else ""
|
507 |
msg_lora = ""
|
508 |
|
509 |
-
print("Config model:", model_name, vae_model, loras_list)
|
510 |
-
|
511 |
## BEGIN MOD
|
|
|
512 |
prompt, neg_prompt = insert_model_recom_prompt(prompt, neg_prompt, model_name)
|
513 |
global lora_model_list
|
514 |
lora_model_list = get_lora_model_list()
|
515 |
## END MOD
|
516 |
|
|
|
|
|
517 |
task = TASK_STABLEPY[task]
|
518 |
|
519 |
params_ip_img = []
|
@@ -536,6 +364,9 @@ class GuiSD:
|
|
536 |
params_ip_mode.append(modeip)
|
537 |
params_ip_scale.append(scaleip)
|
538 |
|
|
|
|
|
|
|
539 |
if task != "txt2img" and not image_control:
|
540 |
raise ValueError("No control image found: To use this function, you have to upload an image in 'Image ControlNet/Inpaint/Img2img'")
|
541 |
|
@@ -665,18 +496,17 @@ class GuiSD:
|
|
665 |
}
|
666 |
|
667 |
self.model.device = torch.device("cuda:0")
|
668 |
-
if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5
|
669 |
self.model.pipe.transformer.to(self.model.device)
|
670 |
print("transformer to cuda")
|
671 |
|
672 |
-
#progress(1, desc="Inference preparation completed. Starting inference...")
|
673 |
-
|
674 |
-
info_state = "" # for yield version
|
675 |
return self.infer_short(self.model, pipe_params), info_state
|
676 |
## END MOD
|
677 |
|
|
|
678 |
def dynamic_gpu_duration(func, duration, *args):
|
679 |
|
|
|
680 |
@spaces.GPU(duration=duration)
|
681 |
def wrapped_func():
|
682 |
return func(*args)
|
@@ -696,7 +526,7 @@ def sd_gen_generate_pipeline(*args):
|
|
696 |
load_lora_cpu = args[-3]
|
697 |
generation_args = args[:-3]
|
698 |
lora_list = [
|
699 |
-
None if item == "None" or item == "" else item
|
700 |
for item in [args[7], args[9], args[11], args[13], args[15]]
|
701 |
]
|
702 |
lora_status = [None] * 5
|
@@ -706,7 +536,7 @@ def sd_gen_generate_pipeline(*args):
|
|
706 |
msg_load_lora = "Updating LoRAs in CPU (Slow but saves GPU usage)..."
|
707 |
|
708 |
#if lora_list != sd_gen.model.lora_memory and lora_list != [None] * 5:
|
709 |
-
# yield
|
710 |
|
711 |
# Load lora in CPU
|
712 |
if load_lora_cpu:
|
@@ -732,14 +562,16 @@ def sd_gen_generate_pipeline(*args):
|
|
732 |
)
|
733 |
gr.Info(f"LoRAs in cache: {lora_cache_msg}")
|
734 |
|
735 |
-
|
|
|
736 |
gr.Info(msg_request)
|
737 |
print(msg_request)
|
738 |
-
|
739 |
-
# yield from sd_gen.generate_pipeline(*generation_args)
|
740 |
|
741 |
start_time = time.time()
|
742 |
|
|
|
|
|
743 |
return dynamic_gpu_duration(
|
744 |
sd_gen.generate_pipeline,
|
745 |
gpu_duration_arg,
|
@@ -747,31 +579,19 @@ def sd_gen_generate_pipeline(*args):
|
|
747 |
)
|
748 |
|
749 |
end_time = time.time()
|
|
|
|
|
|
|
|
|
750 |
|
751 |
if verbose_arg:
|
752 |
-
execution_time = end_time - start_time
|
753 |
-
msg_task_complete = (
|
754 |
-
f"GPU task complete in: {round(execution_time, 0) + 1} seconds"
|
755 |
-
)
|
756 |
gr.Info(msg_task_complete)
|
757 |
print(msg_task_complete)
|
758 |
|
759 |
-
|
760 |
-
if image is None: return ""
|
761 |
|
762 |
-
try:
|
763 |
-
metadata_keys = ['parameters', 'metadata', 'prompt', 'Comment']
|
764 |
-
|
765 |
-
for key in metadata_keys:
|
766 |
-
if key in image.info:
|
767 |
-
return image.info[key]
|
768 |
|
769 |
-
|
770 |
-
|
771 |
-
except Exception as e:
|
772 |
-
return f"Error extracting metadata: {str(e)}"
|
773 |
-
|
774 |
-
@spaces.GPU(duration=20)
|
775 |
def esrgan_upscale(image, upscaler_name, upscaler_size):
|
776 |
if image is None: return None
|
777 |
|
@@ -793,9 +613,11 @@ def esrgan_upscale(image, upscaler_name, upscaler_size):
|
|
793 |
|
794 |
return image_path
|
795 |
|
|
|
796 |
dynamic_gpu_duration.zerogpu = True
|
797 |
sd_gen_generate_pipeline.zerogpu = True
|
798 |
|
|
|
799 |
from pathlib import Path
|
800 |
from PIL import Image
|
801 |
import random, json
|
@@ -1027,14 +849,14 @@ def update_lora_dict(path: str):
|
|
1027 |
def download_lora(dl_urls: str):
|
1028 |
global loras_url_to_path_dict
|
1029 |
dl_path = ""
|
1030 |
-
before = get_local_model_list(
|
1031 |
urls = []
|
1032 |
for url in [url.strip() for url in dl_urls.split(',')]:
|
1033 |
-
local_path = f"{
|
1034 |
if not Path(local_path).exists():
|
1035 |
-
download_things(
|
1036 |
urls.append(url)
|
1037 |
-
after = get_local_model_list(
|
1038 |
new_files = list_sub(after, before)
|
1039 |
i = 0
|
1040 |
for file in new_files:
|
|
|
59 |
logger.setLevel(logging.DEBUG)
|
60 |
|
61 |
from env import (
|
62 |
+
HF_TOKEN, HF_READ_TOKEN, # to use only for private repos
|
63 |
CIVITAI_API_KEY, HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
|
64 |
HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
|
65 |
HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
|
66 |
+
DIRECTORY_MODELS, DIRECTORY_LORAS, DIRECTORY_VAES, DIRECTORY_EMBEDS,
|
67 |
+
DIRECTORY_EMBEDS_SDXL, DIRECTORY_EMBEDS_POSITIVE_SDXL,
|
68 |
+
LOAD_DIFFUSERS_FORMAT_MODEL, DOWNLOAD_MODEL_LIST, DOWNLOAD_LORA_LIST,
|
69 |
+
DOWNLOAD_VAE_LIST, download_embeds)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
70 |
|
71 |
## BEGIN MOD
|
72 |
from modutils import (to_list, list_uniq, list_sub, get_model_id_list, get_tupled_embed_list,
|
73 |
get_tupled_model_list, get_lora_model_list, download_private_repo, download_things)
|
74 |
|
75 |
# - **Download Models**
|
76 |
+
download_model = ", ".join(DOWNLOAD_MODEL_LIST)
|
77 |
# - **Download VAEs**
|
78 |
+
download_vae = ", ".join(DOWNLOAD_VAE_LIST)
|
79 |
# - **Download LoRAs**
|
80 |
+
download_lora = ", ".join(DOWNLOAD_LORA_LIST)
|
81 |
|
82 |
+
#download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO, DIRECTORY_LORAS, True)
|
83 |
+
download_private_repo(HF_VAE_PRIVATE_REPO, DIRECTORY_VAES, False)
|
84 |
|
85 |
+
load_diffusers_format_model = list_uniq(LOAD_DIFFUSERS_FORMAT_MODEL + get_model_id_list())
|
86 |
## END MOD
|
87 |
|
88 |
# Download stuffs
|
89 |
for url in [url.strip() for url in download_model.split(',')]:
|
90 |
if not os.path.exists(f"./models/{url.split('/')[-1]}"):
|
91 |
+
download_things(DIRECTORY_MODELS, url, HF_TOKEN, CIVITAI_API_KEY)
|
92 |
for url in [url.strip() for url in download_vae.split(',')]:
|
93 |
if not os.path.exists(f"./vaes/{url.split('/')[-1]}"):
|
94 |
+
download_things(DIRECTORY_VAES, url, HF_TOKEN, CIVITAI_API_KEY)
|
95 |
for url in [url.strip() for url in download_lora.split(',')]:
|
96 |
if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
|
97 |
+
download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY)
|
98 |
|
99 |
# Download Embeddings
|
100 |
for url_embed in download_embeds:
|
101 |
if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
|
102 |
+
download_things(DIRECTORY_EMBEDS, url_embed, HF_TOKEN, CIVITAI_API_KEY)
|
103 |
|
104 |
# Build list models
|
105 |
+
embed_list = get_model_list(DIRECTORY_EMBEDS)
|
106 |
+
model_list = get_model_list(DIRECTORY_MODELS)
|
107 |
model_list = load_diffusers_format_model + model_list
|
108 |
## BEGIN MOD
|
109 |
lora_model_list = get_lora_model_list()
|
110 |
+
vae_model_list = get_model_list(DIRECTORY_VAES)
|
111 |
vae_model_list.insert(0, "None")
|
112 |
|
113 |
+
#download_private_repo(HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_SDXL, False)
|
114 |
+
#download_private_repo(HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_POSITIVE_SDXL, False)
|
115 |
+
embed_sdxl_list = get_model_list(DIRECTORY_EMBEDS_SDXL) + get_model_list(DIRECTORY_EMBEDS_POSITIVE_SDXL)
|
116 |
|
117 |
def get_embed_list(pipeline_name):
|
118 |
return get_tupled_embed_list(embed_sdxl_list if pipeline_name == "StableDiffusionXLPipeline" else embed_list)
|
|
|
120 |
|
121 |
print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
|
122 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
123 |
## BEGIN MOD
|
124 |
class GuiSD:
|
125 |
+
def __init__(self, stream=True):
|
126 |
self.model = None
|
127 |
+
self.status_loading = False
|
128 |
+
self.sleep_loading = 4
|
129 |
+
self.last_load = datetime.now()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
130 |
|
131 |
def infer_short(self, model, pipe_params, progress=gr.Progress(track_tqdm=True)):
|
132 |
#progress(0, desc="Start inference...")
|
|
|
140 |
return img
|
141 |
|
142 |
def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
|
|
|
|
|
|
|
143 |
vae_model = vae_model if vae_model != "None" else None
|
144 |
model_type = get_model_type(model_name)
|
145 |
+
dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
|
146 |
+
|
147 |
+
if not os.path.exists(model_name):
|
148 |
+
_ = download_diffuser_repo(
|
149 |
+
repo_name=model_name,
|
150 |
+
model_type=model_type,
|
151 |
+
revision="main",
|
152 |
+
token=True,
|
153 |
+
)
|
154 |
+
|
155 |
+
for i in range(68):
|
156 |
+
if not self.status_loading:
|
157 |
+
self.status_loading = True
|
158 |
+
if i > 0:
|
159 |
+
time.sleep(self.sleep_loading)
|
160 |
+
print("Previous model ops...")
|
161 |
+
break
|
162 |
+
time.sleep(0.5)
|
163 |
+
print(f"Waiting queue {i}")
|
164 |
+
yield "Waiting queue"
|
165 |
+
|
166 |
+
self.status_loading = True
|
167 |
+
|
168 |
+
#yield f"Loading model: {model_name}"
|
169 |
|
170 |
if vae_model:
|
171 |
vae_type = "SDXL" if "sdxl" in vae_model.lower() else "SD 1.5"
|
172 |
if model_type != vae_type:
|
173 |
+
gr.Warning(WARNING_MSG_VAE)
|
174 |
+
|
175 |
+
print("Loading model...")
|
176 |
+
|
177 |
+
try:
|
178 |
+
start_time = time.time()
|
179 |
+
|
180 |
+
if self.model is None:
|
181 |
+
self.model = Model_Diffusers(
|
182 |
+
base_model_id=model_name,
|
183 |
+
task_name=TASK_STABLEPY[task],
|
184 |
+
vae_model=vae_model,
|
185 |
+
type_model_precision=dtype_model,
|
186 |
+
retain_task_model_in_cache=False,
|
187 |
+
device="cpu",
|
188 |
+
)
|
189 |
+
else:
|
190 |
+
|
191 |
+
if self.model.base_model_id != model_name:
|
192 |
+
load_now_time = datetime.now()
|
193 |
+
elapsed_time = (load_now_time - self.last_load).total_seconds()
|
194 |
+
|
195 |
+
if elapsed_time <= 8:
|
196 |
+
print("Waiting for the previous model's time ops...")
|
197 |
+
time.sleep(8-elapsed_time)
|
198 |
+
|
199 |
+
self.model.device = torch.device("cpu")
|
200 |
+
self.model.load_pipe(
|
201 |
+
model_name,
|
202 |
+
task_name=TASK_STABLEPY[task],
|
203 |
+
vae_model=vae_model,
|
204 |
+
type_model_precision=dtype_model,
|
205 |
+
retain_task_model_in_cache=False,
|
206 |
+
)
|
207 |
+
|
208 |
+
end_time = time.time()
|
209 |
+
self.sleep_loading = max(min(int(end_time - start_time), 10), 4)
|
210 |
+
except Exception as e:
|
211 |
+
self.last_load = datetime.now()
|
212 |
+
self.status_loading = False
|
213 |
+
self.sleep_loading = 4
|
214 |
+
raise e
|
215 |
+
|
216 |
+
self.last_load = datetime.now()
|
217 |
+
self.status_loading = False
|
218 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
219 |
#yield f"Model loaded: {model_name}"
|
220 |
|
221 |
#@spaces.GPU
|
222 |
+
#@torch.inference_mode()
|
223 |
def generate_pipeline(
|
224 |
self,
|
225 |
prompt,
|
|
|
324 |
mode_ip2,
|
325 |
scale_ip2,
|
326 |
pag_scale,
|
|
|
327 |
):
|
328 |
+
info_state = html_template_message("Navigating latent space...")
|
329 |
+
#yield info_state, gr.update(), gr.update()
|
330 |
+
|
331 |
vae_model = vae_model if vae_model != "None" else None
|
332 |
loras_list = [lora1, lora2, lora3, lora4, lora5]
|
333 |
vae_msg = f"VAE: {vae_model}" if vae_model else ""
|
334 |
msg_lora = ""
|
335 |
|
|
|
|
|
336 |
## BEGIN MOD
|
337 |
+
loras_list = [s if s else "None" for s in loras_list]
|
338 |
prompt, neg_prompt = insert_model_recom_prompt(prompt, neg_prompt, model_name)
|
339 |
global lora_model_list
|
340 |
lora_model_list = get_lora_model_list()
|
341 |
## END MOD
|
342 |
|
343 |
+
print("Config model:", model_name, vae_model, loras_list)
|
344 |
+
|
345 |
task = TASK_STABLEPY[task]
|
346 |
|
347 |
params_ip_img = []
|
|
|
364 |
params_ip_mode.append(modeip)
|
365 |
params_ip_scale.append(scaleip)
|
366 |
|
367 |
+
concurrency = 5
|
368 |
+
self.model.stream_config(concurrency=concurrency, latent_resize_by=1, vae_decoding=False)
|
369 |
+
|
370 |
if task != "txt2img" and not image_control:
|
371 |
raise ValueError("No control image found: To use this function, you have to upload an image in 'Image ControlNet/Inpaint/Img2img'")
|
372 |
|
|
|
496 |
}
|
497 |
|
498 |
self.model.device = torch.device("cuda:0")
|
499 |
+
if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5:
|
500 |
self.model.pipe.transformer.to(self.model.device)
|
501 |
print("transformer to cuda")
|
502 |
|
|
|
|
|
|
|
503 |
return self.infer_short(self.model, pipe_params), info_state
|
504 |
## END MOD
|
505 |
|
506 |
+
|
507 |
def dynamic_gpu_duration(func, duration, *args):
|
508 |
|
509 |
+
@torch.inference_mode()
|
510 |
@spaces.GPU(duration=duration)
|
511 |
def wrapped_func():
|
512 |
return func(*args)
|
|
|
526 |
load_lora_cpu = args[-3]
|
527 |
generation_args = args[:-3]
|
528 |
lora_list = [
|
529 |
+
None if item == "None" or item == "" else item # MOD
|
530 |
for item in [args[7], args[9], args[11], args[13], args[15]]
|
531 |
]
|
532 |
lora_status = [None] * 5
|
|
|
536 |
msg_load_lora = "Updating LoRAs in CPU (Slow but saves GPU usage)..."
|
537 |
|
538 |
#if lora_list != sd_gen.model.lora_memory and lora_list != [None] * 5:
|
539 |
+
# yield msg_load_lora, gr.update(), gr.update()
|
540 |
|
541 |
# Load lora in CPU
|
542 |
if load_lora_cpu:
|
|
|
562 |
)
|
563 |
gr.Info(f"LoRAs in cache: {lora_cache_msg}")
|
564 |
|
565 |
+
msg_request = f"Requesting {gpu_duration_arg}s. of GPU time.\nModel: {sd_gen.model.base_model_id}"
|
566 |
+
if verbose_arg:
|
567 |
gr.Info(msg_request)
|
568 |
print(msg_request)
|
569 |
+
#yield msg_request.replace("\n", "<br>"), gr.update(), gr.update()
|
|
|
570 |
|
571 |
start_time = time.time()
|
572 |
|
573 |
+
# yield from sd_gen.generate_pipeline(*generation_args)
|
574 |
+
#yield from dynamic_gpu_duration(
|
575 |
return dynamic_gpu_duration(
|
576 |
sd_gen.generate_pipeline,
|
577 |
gpu_duration_arg,
|
|
|
579 |
)
|
580 |
|
581 |
end_time = time.time()
|
582 |
+
execution_time = end_time - start_time
|
583 |
+
msg_task_complete = (
|
584 |
+
f"GPU task complete in: {int(round(execution_time, 0) + 1)} seconds"
|
585 |
+
)
|
586 |
|
587 |
if verbose_arg:
|
|
|
|
|
|
|
|
|
588 |
gr.Info(msg_task_complete)
|
589 |
print(msg_task_complete)
|
590 |
|
591 |
+
yield msg_task_complete, gr.update(), gr.update()
|
|
|
592 |
|
|
|
|
|
|
|
|
|
|
|
|
|
593 |
|
594 |
+
@spaces.GPU(duration=15)
|
|
|
|
|
|
|
|
|
|
|
595 |
def esrgan_upscale(image, upscaler_name, upscaler_size):
|
596 |
if image is None: return None
|
597 |
|
|
|
613 |
|
614 |
return image_path
|
615 |
|
616 |
+
|
617 |
dynamic_gpu_duration.zerogpu = True
|
618 |
sd_gen_generate_pipeline.zerogpu = True
|
619 |
|
620 |
+
|
621 |
from pathlib import Path
|
622 |
from PIL import Image
|
623 |
import random, json
|
|
|
849 |
def download_lora(dl_urls: str):
|
850 |
global loras_url_to_path_dict
|
851 |
dl_path = ""
|
852 |
+
before = get_local_model_list(DIRECTORY_LORAS)
|
853 |
urls = []
|
854 |
for url in [url.strip() for url in dl_urls.split(',')]:
|
855 |
+
local_path = f"{DIRECTORY_LORAS}/{url.split('/')[-1]}"
|
856 |
if not Path(local_path).exists():
|
857 |
+
download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY)
|
858 |
urls.append(url)
|
859 |
+
after = get_local_model_list(DIRECTORY_LORAS)
|
860 |
new_files = list_sub(after, before)
|
861 |
i = 0
|
862 |
for file in new_files:
|