import spaces
import os
from stablepy import Model_Diffusers
from stablepy.diffusers_vanilla.model import scheduler_names
from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
import torch
import re
import shutil
import random
import spaces
import gradio as gr
from PIL import Image
import IPython.display
import time, json
from IPython.utils import capture
import logging
from utils.string_utils import extract_parameters
from stablepy import logger
from models.upscaler import upscaler_dict_gui
logging.getLogger("diffusers").setLevel(logging.ERROR)
import diffusers
diffusers.utils.logging.set_verbosity(40)
import warnings
from utils.download_utils import download_things
hf_token: str = os.environ.get("HF_TOKEN")
class GuiSD:
def __init__(self,
model_list,
task_stablepy,
lora_model_list,
embed_list,
stream=True):
self.model = None
print("Loading model...")
self.model = Model_Diffusers(
base_model_id="cagliostrolab/animagine-xl-3.1",
task_name="txt2img",
vae_model=None,
type_model_precision=torch.float16,
retain_task_model_in_cache=False,
)
self.model_list = model_list
self.task_stablepy = task_stablepy
self.lora_model_list = lora_model_list
self.embed_list = embed_list
self.stream = stream
def load_new_model(
self,
model_name,
vae_model,
task,
progress=gr.Progress(track_tqdm=True)):
"""
:param model_name:
:param vae_model:
:param task:
:param progress:
"""
yield f"Loading model: {model_name}"
vae_model = vae_model if vae_model != "None" else None
if model_name in self.model_list:
model_is_xl = "xl" in model_name.lower()
sdxl_in_vae = vae_model and "sdxl" in vae_model.lower()
model_type = "SDXL" if model_is_xl else "SD 1.5"
incompatible_vae = (model_is_xl and vae_model and not sdxl_in_vae) or (not model_is_xl and sdxl_in_vae)
if incompatible_vae:
vae_model = None
self.model.load_pipe(
model_name,
task_name=self.task_stablepy[task],
vae_model=vae_model if vae_model != "None" else None,
type_model_precision=torch.float16,
retain_task_model_in_cache=False,
)
yield f"Model loaded: {model_name}"
@spaces.GPU
def generate_pipeline(
self,
prompt,
neg_prompt,
num_images,
steps,
cfg,
clip_skip,
seed,
lora1,
lora_scale1,
lora2,
lora_scale2,
lora3,
lora_scale3,
lora4,
lora_scale4,
lora5,
lora_scale5,
sampler,
img_height,
img_width,
model_name,
vae_model,
task,
image_control,
preprocessor_name,
preprocess_resolution,
image_resolution,
style_prompt, # list []
style_json_file,
image_mask,
strength,
low_threshold,
high_threshold,
value_threshold,
distance_threshold,
controlnet_output_scaling_in_unet,
controlnet_start_threshold,
controlnet_stop_threshold,
textual_inversion,
syntax_weights,
upscaler_model_path,
upscaler_increases_size,
esrgan_tile,
esrgan_tile_overlap,
hires_steps,
hires_denoising_strength,
hires_sampler,
hires_prompt,
hires_negative_prompt,
hires_before_adetailer,
hires_after_adetailer,
loop_generation,
leave_progress_bar,
disable_progress_bar,
image_previews,
display_images,
save_generated_images,
image_storage_location,
retain_compel_previous_load,
retain_detailfix_model_previous_load,
retain_hires_model_previous_load,
t2i_adapter_preprocessor,
t2i_adapter_conditioning_scale,
t2i_adapter_conditioning_factor,
xformers_memory_efficient_attention,
freeu,
generator_in_cpu,
adetailer_inpaint_only,
adetailer_verbose,
adetailer_sampler,
adetailer_active_a,
prompt_ad_a,
negative_prompt_ad_a,
strength_ad_a,
face_detector_ad_a,
person_detector_ad_a,
hand_detector_ad_a,
mask_dilation_a,
mask_blur_a,
mask_padding_a,
adetailer_active_b,
prompt_ad_b,
negative_prompt_ad_b,
strength_ad_b,
face_detector_ad_b,
person_detector_ad_b,
hand_detector_ad_b,
mask_dilation_b,
mask_blur_b,
mask_padding_b,
retain_task_cache_gui,
image_ip1,
mask_ip1,
model_ip1,
mode_ip1,
scale_ip1,
image_ip2,
mask_ip2,
model_ip2,
mode_ip2,
scale_ip2):
vae_model = vae_model if vae_model != "None" else None
loras_list: list = [lora1, lora2, lora3, lora4, lora5]
vae_msg: str = f"VAE: {vae_model}" if vae_model else ""
msg_lora: list = []
if model_name in self.model_list:
model_is_xl = "xl" in model_name.lower()
sdxl_in_vae = vae_model and "sdxl" in vae_model.lower()
model_type = "SDXL" if model_is_xl else "SD 1.5"
incompatible_vae = ((model_is_xl and
vae_model and
not sdxl_in_vae) or
(not model_is_xl and
sdxl_in_vae))
if incompatible_vae:
msg_inc_vae = (
f"The selected VAE is for a {'SD 1.5' if model_is_xl else 'SDXL'} model, but you"
f" are using a {model_type} model. The default VAE "
"will be used."
)
gr.Info(msg_inc_vae)
vae_msg = msg_inc_vae
vae_model = None
for la in loras_list:
if la is None or la == "None" or la not in self.lora_model_list:
continue
print(la)
lora_type = ("animetarot" in la.lower() or "Hyper-SD15-8steps".lower() in la.lower())
if (model_is_xl and lora_type) or (not model_is_xl and not lora_type):
msg_inc_lora = f"The LoRA {la} is for {'SD 1.5' if model_is_xl else 'SDXL'}, but you are using {model_type}."
gr.Info(msg_inc_lora)
msg_lora.append(msg_inc_lora)
task = self.task_stablepy[task]
params_ip_img: list = []
params_ip_msk: list = []
params_ip_model: list = []
params_ip_mode: list = []
params_ip_scale: list = []
all_adapters = [
(image_ip1, mask_ip1, model_ip1, mode_ip1, scale_ip1),
(image_ip2, mask_ip2, model_ip2, mode_ip2, scale_ip2),
]
for (imgip,
mskip,
modelip,
modeip,
scaleip) in all_adapters:
if imgip:
params_ip_img.append(imgip)
if mskip:
params_ip_msk.append(mskip)
params_ip_model.append(modelip)
params_ip_mode.append(modeip)
params_ip_scale.append(scaleip)
# First load
model_precision = torch.float16
if not self.model:
from modelstream import Model_Diffusers2
print("Loading model...")
self.model = Model_Diffusers2(
base_model_id=model_name,
task_name=task,
vae_model=vae_model if vae_model != "None" else None,
type_model_precision=model_precision,
retain_task_model_in_cache=retain_task_cache_gui,
)
if task != "txt2img" and not image_control:
raise ValueError(
"No control image found: To use this function, "
"you have to upload an image in 'Image ControlNet/Inpaint/Img2img'"
)
if task == "inpaint" and not image_mask:
raise ValueError("No mask image found: Specify one in 'Image Mask'")
if upscaler_model_path in [None, "Lanczos", "Nearest"]:
upscaler_model = upscaler_model_path
else:
directory_upscalers = 'upscalers'
os.makedirs(directory_upscalers, exist_ok=True)
url_upscaler = upscaler_dict_gui[upscaler_model_path]
if not os.path.exists(f"./upscalers/{url_upscaler.split('/')[-1]}"):
download_things(
directory_upscalers,
url_upscaler,
hf_token
)
upscaler_model = f"./upscalers/{url_upscaler.split('/')[-1]}"
logging.getLogger("ultralytics").setLevel(logging.INFO if adetailer_verbose else logging.ERROR)
print("Config model:", model_name, vae_model, loras_list)
self.model.load_pipe(
model_name,
task_name=task,
vae_model=vae_model if vae_model != "None" else None,
type_model_precision=model_precision,
retain_task_model_in_cache=retain_task_cache_gui,
)
if textual_inversion and self.model.class_name == "StableDiffusionXLPipeline":
print("No Textual inversion for SDXL")
adetailer_params_A: dict = {
"face_detector_ad": face_detector_ad_a,
"person_detector_ad": person_detector_ad_a,
"hand_detector_ad": hand_detector_ad_a,
"prompt": prompt_ad_a,
"negative_prompt": negative_prompt_ad_a,
"strength": strength_ad_a,
# "image_list_task" : None,
"mask_dilation": mask_dilation_a,
"mask_blur": mask_blur_a,
"mask_padding": mask_padding_a,
"inpaint_only": adetailer_inpaint_only,
"sampler": adetailer_sampler,
}
adetailer_params_B: dict = {
"face_detector_ad": face_detector_ad_b,
"person_detector_ad": person_detector_ad_b,
"hand_detector_ad": hand_detector_ad_b,
"prompt": prompt_ad_b,
"negative_prompt": negative_prompt_ad_b,
"strength": strength_ad_b,
# "image_list_task" : None,
"mask_dilation": mask_dilation_b,
"mask_blur": mask_blur_b,
"mask_padding": mask_padding_b,
}
pipe_params: dict = {
"prompt": prompt,
"negative_prompt": neg_prompt,
"img_height": img_height,
"img_width": img_width,
"num_images": num_images,
"num_steps": steps,
"guidance_scale": cfg,
"clip_skip": clip_skip,
"seed": seed,
"image": image_control,
"preprocessor_name": preprocessor_name,
"preprocess_resolution": preprocess_resolution,
"image_resolution": image_resolution,
"style_prompt": style_prompt if style_prompt else "",
"style_json_file": "",
"image_mask": image_mask, # only for Inpaint
"strength": strength, # only for Inpaint or ...
"low_threshold": low_threshold,
"high_threshold": high_threshold,
"value_threshold": value_threshold,
"distance_threshold": distance_threshold,
"lora_A": lora1 if lora1 != "None" else None,
"lora_scale_A": lora_scale1,
"lora_B": lora2 if lora2 != "None" else None,
"lora_scale_B": lora_scale2,
"lora_C": lora3 if lora3 != "None" else None,
"lora_scale_C": lora_scale3,
"lora_D": lora4 if lora4 != "None" else None,
"lora_scale_D": lora_scale4,
"lora_E": lora5 if lora5 != "None" else None,
"lora_scale_E": lora_scale5,
"textual_inversion": self.embed_list if textual_inversion and self.model.class_name != "StableDiffusionXLPipeline" else [],
"syntax_weights": syntax_weights, # "Classic"
"sampler": sampler,
"xformers_memory_efficient_attention": xformers_memory_efficient_attention,
"gui_active": True,
"loop_generation": loop_generation,
"controlnet_conditioning_scale": float(controlnet_output_scaling_in_unet),
"control_guidance_start": float(controlnet_start_threshold),
"control_guidance_end": float(controlnet_stop_threshold),
"generator_in_cpu": generator_in_cpu,
"FreeU": freeu,
"adetailer_A": adetailer_active_a,
"adetailer_A_params": adetailer_params_A,
"adetailer_B": adetailer_active_b,
"adetailer_B_params": adetailer_params_B,
"leave_progress_bar": leave_progress_bar,
"disable_progress_bar": disable_progress_bar,
"image_previews": image_previews,
"display_images": display_images,
"save_generated_images": save_generated_images,
"image_storage_location": image_storage_location,
"retain_compel_previous_load": retain_compel_previous_load,
"retain_detailfix_model_previous_load": retain_detailfix_model_previous_load,
"retain_hires_model_previous_load": retain_hires_model_previous_load,
"t2i_adapter_preprocessor": t2i_adapter_preprocessor,
"t2i_adapter_conditioning_scale": float(t2i_adapter_conditioning_scale),
"t2i_adapter_conditioning_factor": float(t2i_adapter_conditioning_factor),
"upscaler_model_path": upscaler_model,
"upscaler_increases_size": upscaler_increases_size,
"esrgan_tile": esrgan_tile,
"esrgan_tile_overlap": esrgan_tile_overlap,
"hires_steps": hires_steps,
"hires_denoising_strength": hires_denoising_strength,
"hires_prompt": hires_prompt,
"hires_negative_prompt": hires_negative_prompt,
"hires_sampler": hires_sampler,
"hires_before_adetailer": hires_before_adetailer,
"hires_after_adetailer": hires_after_adetailer,
"ip_adapter_image": params_ip_img,
"ip_adapter_mask": params_ip_msk,
"ip_adapter_model": params_ip_model,
"ip_adapter_mode": params_ip_mode,
"ip_adapter_scale": params_ip_scale,
}
# print(pipe_params)
random_number = random.randint(1, 100)
if random_number < 25 and num_images < 3:
if (not upscaler_model and
steps < 45 and
task in ["txt2img", "img2img"] and
not adetailer_active_a and
not adetailer_active_b):
num_images *= 2
pipe_params["num_images"] = num_images
gr.Info("Num images x 2 🎉")
# Maybe fix lora issue: 'Cannot copy out of meta tensor; no data!''
self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
info_state = f"PROCESSING"
for img, seed, data in self.model(**pipe_params):
info_state += "."
if data:
info_state = f"COMPLETED. Seeds: {str(seed)}"
if vae_msg:
info_state = info_state + "
" + vae_msg
if msg_lora:
info_state = info_state + "
" + "
".join(msg_lora)
yield img, info_state