Upload dc.py

dc.py

@@ -1,33 +1,52 @@
 import spaces
 import os
 from stablepy import Model_Diffusers
+from constants import (
+    PREPROCESSOR_CONTROLNET,
+    TASK_STABLEPY,
+    TASK_MODEL_LIST,
+    UPSCALER_DICT_GUI,
+    UPSCALER_KEYS,
+    PROMPT_W_OPTIONS,
+    WARNING_MSG_VAE,
+    SDXL_TASK,
+    MODEL_TYPE_TASK,
+    POST_PROCESSING_SAMPLER,
+
+)
 from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
-from stablepy.diffusers_vanilla.constants import FLUX_CN_UNION_MODES
 import torch
 import re
-from huggingface_hub import HfApi
 from stablepy import (
-    CONTROLNET_MODEL_IDS,
-    VALID_TASKS,
-    T2I_PREPROCESSOR_NAME,
-    FLASH_LORA,
-    SCHEDULER_CONFIG_MAP,
     scheduler_names,
-    IP_ADAPTER_MODELS,
     IP_ADAPTERS_SD,
     IP_ADAPTERS_SDXL,
-    REPO_IMAGE_ENCODER,
-    ALL_PROMPT_WEIGHT_OPTIONS,
-    SD15_TASKS,
-    SDXL_TASKS,
 )
 import time
 from PIL import ImageFile
-#import urllib.parse
+from utils import (
+    get_model_list,
+    extract_parameters,
+    get_model_type,
+    extract_exif_data,
+    create_mask_now,
+    download_diffuser_repo,
+    progress_step_bar,
+    html_template_message,
+)
+from datetime import datetime
+import gradio as gr
+import logging
+import diffusers
+import warnings
+from stablepy import logger
+# import urllib.parse
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
+# os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"
 print(os.getenv("SPACES_ZERO_GPU"))
 
+## BEGIN MOD
 import gradio as gr
 import logging
 logging.getLogger("diffusers").setLevel(logging.ERROR)
@@ -38,205 +57,63 @@ warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffuse
 warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
 warnings.filterwarnings(action="ignore", category=FutureWarning, module="transformers")
 from stablepy import logger
-logger.setLevel(logging.CRITICAL)
+logger.setLevel(logging.DEBUG)
 
 from env import (
-    HF_TOKEN, hf_read_token, # to use only for private repos
+    HF_TOKEN, HF_READ_TOKEN, # to use only for private repos
     CIVITAI_API_KEY, HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
     HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
     HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
-    directory_models, directory_loras, directory_vaes, directory_embeds,
-    directory_embeds_sdxl, directory_embeds_positive_sdxl,
-    load_diffusers_format_model, download_model_list, download_lora_list,
-    download_vae_list, download_embeds)
-
-PREPROCESSOR_CONTROLNET = {
-  "openpose": [
-    "Openpose",
-    "None",
-  ],
-  "scribble": [
-    "HED",
-    "PidiNet",
-    "None",
-  ],
-  "softedge": [
-    "PidiNet",
-    "HED",
-    "HED safe",
-    "PidiNet safe",
-    "None",
-  ],
-  "segmentation": [
-    "UPerNet",
-    "None",
-  ],
-  "depth": [
-    "DPT",
-    "Midas",
-    "None",
-  ],
-  "normalbae": [
-    "NormalBae",
-    "None",
-  ],
-  "lineart": [
-    "Lineart",
-    "Lineart coarse",
-    "Lineart (anime)",
-    "None",
-    "None (anime)",
-  ],
-  "lineart_anime": [
-    "Lineart",
-    "Lineart coarse",
-    "Lineart (anime)",
-    "None",
-    "None (anime)",
-  ],
-  "shuffle": [
-    "ContentShuffle",
-    "None",
-  ],
-  "canny": [
-    "Canny",
-    "None",
-  ],
-  "mlsd": [
-    "MLSD",
-    "None",
-  ],
-  "ip2p": [
-    "ip2p"
-  ],
-  "recolor": [
-    "Recolor luminance",
-    "Recolor intensity",
-    "None",
-  ],
-  "tile": [
-    "Mild Blur",
-    "Moderate Blur",
-    "Heavy Blur",
-    "None",
-  ],
-}
-
-TASK_STABLEPY = {
-    'txt2img': 'txt2img',
-    'img2img': 'img2img',
-    'inpaint': 'inpaint',
-    # 'canny T2I Adapter': 'sdxl_canny_t2i',  # NO HAVE STEP CALLBACK PARAMETERS SO NOT WORKS WITH DIFFUSERS 0.29.0
-    # 'sketch  T2I Adapter': 'sdxl_sketch_t2i',
-    # 'lineart  T2I Adapter': 'sdxl_lineart_t2i',
-    # 'depth-midas  T2I Adapter': 'sdxl_depth-midas_t2i',
-    # 'openpose  T2I Adapter': 'sdxl_openpose_t2i',
-    'openpose ControlNet': 'openpose',
-    'canny ControlNet': 'canny',
-    'mlsd ControlNet': 'mlsd',
-    'scribble ControlNet': 'scribble',
-    'softedge ControlNet': 'softedge',
-    'segmentation ControlNet': 'segmentation',
-    'depth ControlNet': 'depth',
-    'normalbae ControlNet': 'normalbae',
-    'lineart ControlNet': 'lineart',
-    'lineart_anime ControlNet': 'lineart_anime',
-    'shuffle ControlNet': 'shuffle',
-    'ip2p ControlNet': 'ip2p',
-    'optical pattern ControlNet': 'pattern',
-    'recolor ControlNet': 'recolor',
-    'tile ControlNet': 'tile',
-}
-
-TASK_MODEL_LIST = list(TASK_STABLEPY.keys())
-
-UPSCALER_DICT_GUI = {
-    None: None,
-    "Lanczos": "Lanczos",
-    "Nearest": "Nearest",
-    'Latent': 'Latent',
-    'Latent (antialiased)': 'Latent (antialiased)',
-    'Latent (bicubic)': 'Latent (bicubic)',
-    'Latent (bicubic antialiased)': 'Latent (bicubic antialiased)',
-    'Latent (nearest)': 'Latent (nearest)',
-    'Latent (nearest-exact)': 'Latent (nearest-exact)',
-    "RealESRGAN_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
-    "RealESRNet_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
-    "RealESRGAN_x4plus_anime_6B": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
-    "RealESRGAN_x2plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
-    "realesr-animevideov3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
-    "realesr-general-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
-    "realesr-general-wdn-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
-    "4x-UltraSharp": "https://huggingface.co/Shandypur/ESRGAN-4x-UltraSharp/resolve/main/4x-UltraSharp.pth",
-    "4x_foolhardy_Remacri": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
-    "Remacri4xExtraSmoother": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/Remacri%204x%20ExtraSmoother.pth",
-    "AnimeSharp4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/AnimeSharp%204x.pth",
-    "lollypop": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/lollypop.pth",
-    "RealisticRescaler4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/RealisticRescaler%204x.pth",
-    "NickelbackFS4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/NickelbackFS%204x.pth"
-}
-
-UPSCALER_KEYS = list(UPSCALER_DICT_GUI.keys())
-
-
-def get_model_list(directory_path):
-    model_list = []
-    valid_extensions = {'.ckpt', '.pt', '.pth', '.safetensors', '.bin'}
-
-    for filename in os.listdir(directory_path):
-        if os.path.splitext(filename)[1] in valid_extensions:
-            # name_without_extension = os.path.splitext(filename)[0]
-            file_path = os.path.join(directory_path, filename)
-            # model_list.append((name_without_extension, file_path))
-            model_list.append(file_path)
-            print('\033[34mFILE: ' + file_path + '\033[0m')
-    return model_list
+    DIRECTORY_MODELS, DIRECTORY_LORAS, DIRECTORY_VAES, DIRECTORY_EMBEDS,
+    DIRECTORY_EMBEDS_SDXL, DIRECTORY_EMBEDS_POSITIVE_SDXL,
+    LOAD_DIFFUSERS_FORMAT_MODEL, DOWNLOAD_MODEL_LIST, DOWNLOAD_LORA_LIST,
+    DOWNLOAD_VAE_LIST, DOWNLOAD_EMBEDS)
 
-## BEGIN MOD
 from modutils import (to_list, list_uniq, list_sub, get_model_id_list, get_tupled_embed_list,
                       get_tupled_model_list, get_lora_model_list, download_private_repo, download_things)
 
 # - **Download Models**
-download_model = ", ".join(download_model_list)
+download_model = ", ".join(DOWNLOAD_MODEL_LIST)
 # - **Download VAEs**
-download_vae = ", ".join(download_vae_list)
+download_vae = ", ".join(DOWNLOAD_VAE_LIST)
 # - **Download LoRAs**
-download_lora = ", ".join(download_lora_list)
+download_lora = ", ".join(DOWNLOAD_LORA_LIST)
 
-#download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO, directory_loras, True)
-download_private_repo(HF_VAE_PRIVATE_REPO, directory_vaes, False)
+#download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO, DIRECTORY_LORAS, True)
+download_private_repo(HF_VAE_PRIVATE_REPO, DIRECTORY_VAES, False)
 
-load_diffusers_format_model = list_uniq(load_diffusers_format_model + get_model_id_list())
+load_diffusers_format_model = list_uniq(LOAD_DIFFUSERS_FORMAT_MODEL + get_model_id_list())
 ## END MOD
 
 # Download stuffs
 for url in [url.strip() for url in download_model.split(',')]:
     if not os.path.exists(f"./models/{url.split('/')[-1]}"):
-        download_things(directory_models, url, HF_TOKEN, CIVITAI_API_KEY)
+        download_things(DIRECTORY_MODELS, url, HF_TOKEN, CIVITAI_API_KEY)
 for url in [url.strip() for url in download_vae.split(',')]:
     if not os.path.exists(f"./vaes/{url.split('/')[-1]}"):
-        download_things(directory_vaes, url, HF_TOKEN, CIVITAI_API_KEY)
+        download_things(DIRECTORY_VAES, url, HF_TOKEN, CIVITAI_API_KEY)
 for url in [url.strip() for url in download_lora.split(',')]:
     if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
-        download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
+        download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY)
 
 # Download Embeddings
-for url_embed in download_embeds:
+for url_embed in DOWNLOAD_EMBEDS:
     if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
-        download_things(directory_embeds, url_embed, HF_TOKEN, CIVITAI_API_KEY)
+        download_things(DIRECTORY_EMBEDS, url_embed, HF_TOKEN, CIVITAI_API_KEY)
 
 # Build list models
-embed_list = get_model_list(directory_embeds)
-model_list = get_model_list(directory_models)
+embed_list = get_model_list(DIRECTORY_EMBEDS)
+model_list = get_model_list(DIRECTORY_MODELS)
 model_list = load_diffusers_format_model + model_list
+
 ## BEGIN MOD
 lora_model_list = get_lora_model_list()
-vae_model_list = get_model_list(directory_vaes)
+vae_model_list = get_model_list(DIRECTORY_VAES)
 vae_model_list.insert(0, "None")
 
-#download_private_repo(HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, directory_embeds_sdxl, False)
-#download_private_repo(HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO, directory_embeds_positive_sdxl, False)
-embed_sdxl_list = get_model_list(directory_embeds_sdxl) + get_model_list(directory_embeds_positive_sdxl)
+#download_private_repo(HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_SDXL, False)
+#download_private_repo(HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_POSITIVE_SDXL, False)
+embed_sdxl_list = get_model_list(DIRECTORY_EMBEDS_SDXL) + get_model_list(DIRECTORY_EMBEDS_POSITIVE_SDXL)
 
 def get_embed_list(pipeline_name):
     return get_tupled_embed_list(embed_sdxl_list if pipeline_name == "StableDiffusionXLPipeline" else embed_list)
@@ -244,99 +121,13 @@ def get_embed_list(pipeline_name):
 
 print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
 
-msg_inc_vae = (
-    "Use the right VAE for your model to maintain image quality. The wrong"
-    " VAE can lead to poor results, like blurriness in the generated images."
-)
-
-SDXL_TASK = [k for k, v in TASK_STABLEPY.items() if v in SDXL_TASKS]
-SD_TASK = [k for k, v in TASK_STABLEPY.items() if v in SD15_TASKS]
-FLUX_TASK = list(TASK_STABLEPY.keys())[:3] + [k for k, v in TASK_STABLEPY.items() if v in FLUX_CN_UNION_MODES.keys()]
-
-MODEL_TYPE_TASK = {
-    "SD 1.5": SD_TASK,
-    "SDXL": SDXL_TASK,
-    "FLUX": FLUX_TASK,
-}
-
-MODEL_TYPE_CLASS = {
-    "diffusers:StableDiffusionPipeline": "SD 1.5",
-    "diffusers:StableDiffusionXLPipeline": "SDXL",
-    "diffusers:FluxPipeline": "FLUX",
-}
-
-POST_PROCESSING_SAMPLER = ["Use same sampler"] + scheduler_names[:-2]
-
-def extract_parameters(input_string):
-    parameters = {}
-    input_string = input_string.replace("\n", "")
-
-    if "Negative prompt:" not in input_string:
-        if "Steps:" in input_string:
-            input_string = input_string.replace("Steps:", "Negative prompt: Steps:")
-        else:
-            print("Invalid metadata")
-            parameters["prompt"] = input_string
-            return parameters
-
-    parm = input_string.split("Negative prompt:")
-    parameters["prompt"] = parm[0].strip()
-    if "Steps:" not in parm[1]:
-        print("Steps not detected")
-        parameters["neg_prompt"] = parm[1].strip()
-        return parameters
-    parm = parm[1].split("Steps:")
-    parameters["neg_prompt"] = parm[0].strip()
-    input_string = "Steps:" + parm[1]
-
-    # Extracting Steps
-    steps_match = re.search(r'Steps: (\d+)', input_string)
-    if steps_match:
-        parameters['Steps'] = int(steps_match.group(1))
-
-    # Extracting Size
-    size_match = re.search(r'Size: (\d+x\d+)', input_string)
-    if size_match:
-        parameters['Size'] = size_match.group(1)
-        width, height = map(int, parameters['Size'].split('x'))
-        parameters['width'] = width
-        parameters['height'] = height
-
-    # Extracting other parameters
-    other_parameters = re.findall(r'(\w+): (.*?)(?=, \w+|$)', input_string)
-    for param in other_parameters:
-        parameters[param[0]] = param[1].strip('"')
-
-    return parameters
-
-def get_model_type(repo_id: str):
-    api = HfApi(token=os.environ.get("HF_TOKEN"))  # if use private or gated model
-    default = "SD 1.5"
-    try:
-        model = api.model_info(repo_id=repo_id, timeout=5.0)
-        tags = model.tags
-        for tag in tags:
-            if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
-    except Exception:
-        return default
-    return default
-
 ## BEGIN MOD
 class GuiSD:
-    def __init__(self):
+    def __init__(self, stream=True):
         self.model = None
-    
-        print("Loading model...")
-        self.model = Model_Diffusers(
-            base_model_id="Lykon/dreamshaper-8",
-            task_name="txt2img",
-            vae_model=None,
-            type_model_precision=torch.float16,
-            retain_task_model_in_cache=False,
-            device="cpu",
-        )
-        self.model.load_beta_styles()
-        #self.model.device = torch.device("cpu") #
+        self.status_loading = False
+        self.sleep_loading = 4
+        self.last_load = datetime.now()
 
     def infer_short(self, model, pipe_params, progress=gr.Progress(track_tqdm=True)):
         #progress(0, desc="Start inference...")
@@ -350,31 +141,86 @@ class GuiSD:
         return img
 
     def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
-
-        #yield f"Loading model: {model_name}"
-        
         vae_model = vae_model if vae_model != "None" else None
         model_type = get_model_type(model_name)
+        dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
+
+        if not os.path.exists(model_name):
+            _ = download_diffuser_repo(
+                repo_name=model_name,
+                model_type=model_type,
+                revision="main",
+                token=True,
+            )
+
+        '''for i in range(68):
+            if not self.status_loading:
+                self.status_loading = True
+                if i > 0:
+                    time.sleep(self.sleep_loading)
+                    print("Previous model ops...")
+                break
+            time.sleep(0.5)
+            print(f"Waiting queue {i}")
+            yield "Waiting queue"
+
+        self.status_loading = True
+
+        yield f"Loading model: {model_name}"'''
 
         if vae_model:
             vae_type = "SDXL" if "sdxl" in vae_model.lower() else "SD 1.5"
             if model_type != vae_type:
-                gr.Warning(msg_inc_vae)
+                gr.Warning(WARNING_MSG_VAE)
+
+        print("Loading model...")
+
+        try:
+            start_time = time.time()
+
+            if self.model is None:
+                self.model = Model_Diffusers(
+                    base_model_id=model_name,
+                    task_name=TASK_STABLEPY[task],
+                    vae_model=vae_model,
+                    type_model_precision=dtype_model,
+                    retain_task_model_in_cache=False,
+                    device="cpu",
+                )
+            else:
+
+                if self.model.base_model_id != model_name:
+                    load_now_time = datetime.now()
+                    elapsed_time = (load_now_time - self.last_load).total_seconds()
+
+                    if elapsed_time <= 8:
+                        print("Waiting for the previous model's time ops...")
+                        time.sleep(8-elapsed_time)
+
+                self.model.device = torch.device("cpu")
+                self.model.load_pipe(
+                    model_name,
+                    task_name=TASK_STABLEPY[task],
+                    vae_model=vae_model,
+                    type_model_precision=dtype_model,
+                    retain_task_model_in_cache=False,
+                )
+
+            end_time = time.time()
+            self.sleep_loading = max(min(int(end_time - start_time), 10), 4)
+        except Exception as e:
+            self.last_load = datetime.now()
+            self.status_loading = False
+            self.sleep_loading = 4
+            raise e
+
+        self.last_load = datetime.now()
+        self.status_loading = False
 
-        self.model.device = torch.device("cpu")
-        dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
-        
-        self.model.load_pipe(
-            model_name,
-            task_name=TASK_STABLEPY[task],
-            vae_model=vae_model if vae_model != "None" else None,
-            type_model_precision=dtype_model,
-            retain_task_model_in_cache=False,
-        )
         #yield f"Model loaded: {model_name}"
 
     #@spaces.GPU
-    @torch.inference_mode()
+    #@torch.inference_mode()
     def generate_pipeline(
         self,
         prompt,
@@ -479,23 +325,24 @@ class GuiSD:
         mode_ip2,
         scale_ip2,
         pag_scale,
-        #progress=gr.Progress(track_tqdm=True),
     ):
-        #progress(0, desc="Preparing inference...")
-        
+        info_state = html_template_message("Navigating latent space...")
+        #yield info_state, gr.update(), gr.update()
+
         vae_model = vae_model if vae_model != "None" else None
         loras_list = [lora1, lora2, lora3, lora4, lora5]
         vae_msg = f"VAE: {vae_model}" if vae_model else ""
         msg_lora = ""
 
-        print("Config model:", model_name, vae_model, loras_list)
-
 ## BEGIN MOD
+        loras_list = [s if s else "None" for s in loras_list]
         prompt, neg_prompt = insert_model_recom_prompt(prompt, neg_prompt, model_name)
         global lora_model_list
         lora_model_list = get_lora_model_list()
 ## END MOD
         
+        print("Config model:", model_name, vae_model, loras_list)
+
         task = TASK_STABLEPY[task]
 
         params_ip_img = []
@@ -518,6 +365,9 @@ class GuiSD:
                 params_ip_mode.append(modeip)
                 params_ip_scale.append(scaleip)
 
+        concurrency = 5
+        self.model.stream_config(concurrency=concurrency, latent_resize_by=1, vae_decoding=False)
+
         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'")
 
@@ -589,15 +439,15 @@ class GuiSD:
             "high_threshold": high_threshold,
             "value_threshold": value_threshold,
             "distance_threshold": distance_threshold,
-            "lora_A": lora1 if lora1 != "None" and lora1 != "" else None,
+            "lora_A": lora1 if lora1 != "None" else None,
             "lora_scale_A": lora_scale1,
-            "lora_B": lora2 if lora2 != "None" and lora2 != "" else None,
+            "lora_B": lora2 if lora2 != "None" else None,
             "lora_scale_B": lora_scale2,
-            "lora_C": lora3 if lora3 != "None" and lora3 != "" else None,
+            "lora_C": lora3 if lora3 != "None" else None,
             "lora_scale_C": lora_scale3,
-            "lora_D": lora4 if lora4 != "None" and lora4 != "" else None,
+            "lora_D": lora4 if lora4 != "None" else None,
             "lora_scale_D": lora_scale4,
-            "lora_E": lora5 if lora5 != "None" and lora5 != "" else None,
+            "lora_E": lora5 if lora5 != "None" else None,
             "lora_scale_E": lora_scale5,
 ## BEGIN MOD
             "textual_inversion": get_embed_list(self.model.class_name) if textual_inversion else [],
@@ -647,18 +497,59 @@ class GuiSD:
         }
 
         self.model.device = torch.device("cuda:0")
-        if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5 and loras_list != [""] * 5:
+        if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5:
             self.model.pipe.transformer.to(self.model.device)
             print("transformer to cuda")
 
-        #progress(1, desc="Inference preparation completed. Starting inference...")
+        #return self.infer_short(self.model, pipe_params), info_state
+
+        actual_progress = 0
+        info_images = gr.update()
+        for img, seed, image_path, metadata in self.model(**pipe_params):
+            info_state = progress_step_bar(actual_progress, steps)
+            actual_progress += concurrency
+            if image_path:
+                info_images = f"Seeds: {str(seed)}"
+                if vae_msg:
+                    info_images = info_images + "<br>" + vae_msg
+
+                if "Cannot copy out of meta tensor; no data!" in self.model.last_lora_error:
+                    msg_ram = "Unable to process the LoRAs due to high RAM usage; please try again later."
+                    print(msg_ram)
+                    msg_lora += f"<br>{msg_ram}"
+
+                for status, lora in zip(self.model.lora_status, self.model.lora_memory):
+                    if status:
+                        msg_lora += f"<br>Loaded: {lora}"
+                    elif status is not None:
+                        msg_lora += f"<br>Error with: {lora}"
+
+                if msg_lora:
+                    info_images += msg_lora
+
+                info_images = info_images + "<br>" + "GENERATION DATA:<br>" + metadata[0].replace("\n", "<br>") + "<br>-------<br>"
+
+                download_links = "<br>".join(
+                    [
+                        f'<a href="{path.replace("/images/", "/file=/home/user/app/images/")}" download="{os.path.basename(path)}">Download Image {i + 1}</a>'
+                        for i, path in enumerate(image_path)
+                    ]
+                )
+                if save_generated_images:
+                    info_images += f"<br>{download_links}"
 
-        info_state = "" # for yield version
-        return self.infer_short(self.model, pipe_params), info_state
+## BEGIN MOD
+                img = save_images(img, metadata)
 ## END MOD
 
+                info_state = "COMPLETE"
+
+            #yield info_state, img, info_images
+            return info_state, img, info_images
+
 def dynamic_gpu_duration(func, duration, *args):
 
+    @torch.inference_mode()
     @spaces.GPU(duration=duration)
     def wrapped_func():
         return func(*args)
@@ -678,7 +569,7 @@ def sd_gen_generate_pipeline(*args):
     load_lora_cpu = args[-3]
     generation_args = args[:-3]
     lora_list = [
-        None if item == "None" or item == "" else item
+        None if item == "None" or item == "" else item # MOD
         for item in [args[7], args[9], args[11], args[13], args[15]]
     ]
     lora_status = [None] * 5
@@ -688,7 +579,7 @@ def sd_gen_generate_pipeline(*args):
         msg_load_lora = "Updating LoRAs in CPU (Slow but saves GPU usage)..."
 
     #if lora_list != sd_gen.model.lora_memory and lora_list != [None] * 5:
-    #    yield None, msg_load_lora
+    #    yield msg_load_lora, gr.update(), gr.update()
 
     # Load lora in CPU
     if load_lora_cpu:
@@ -714,14 +605,16 @@ def sd_gen_generate_pipeline(*args):
             )
             gr.Info(f"LoRAs in cache: {lora_cache_msg}")
 
-        msg_request = f"Requesting {gpu_duration_arg}s. of GPU time"
+    msg_request = f"Requesting {gpu_duration_arg}s. of GPU time.\nModel: {sd_gen.model.base_model_id}"
+    if verbose_arg:
         gr.Info(msg_request)
         print(msg_request)
-
-    # yield from sd_gen.generate_pipeline(*generation_args)
+    #yield msg_request.replace("\n", "<br>"), gr.update(), gr.update()
 
     start_time = time.time()
 
+    # yield from sd_gen.generate_pipeline(*generation_args)
+    #yield from dynamic_gpu_duration(
     return dynamic_gpu_duration(
         sd_gen.generate_pipeline,
         gpu_duration_arg,
@@ -729,31 +622,19 @@ def sd_gen_generate_pipeline(*args):
     )
 
     end_time = time.time()
+    execution_time = end_time - start_time
+    msg_task_complete = (
+        f"GPU task complete in: {int(round(execution_time, 0) + 1)} seconds"
+    )
 
     if verbose_arg:
-        execution_time = end_time - start_time
-        msg_task_complete = (
-            f"GPU task complete in: {round(execution_time, 0) + 1} seconds"
-        )
         gr.Info(msg_task_complete)
         print(msg_task_complete)
 
-def extract_exif_data(image):
-    if image is None: return ""
-
-    try:
-        metadata_keys = ['parameters', 'metadata', 'prompt', 'Comment']
+    yield msg_task_complete, gr.update(), gr.update()
 
-        for key in metadata_keys:
-            if key in image.info:
-                return image.info[key]
 
-        return str(image.info)
-
-    except Exception as e:
-        return f"Error extracting metadata: {str(e)}"
-
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=15)
 def esrgan_upscale(image, upscaler_name, upscaler_size):
     if image is None: return None
 
@@ -775,18 +656,21 @@ def esrgan_upscale(image, upscaler_name, upscaler_size):
 
     return image_path
 
+
 dynamic_gpu_duration.zerogpu = True
 sd_gen_generate_pipeline.zerogpu = True
+sd_gen = GuiSD()
+
 
 from pathlib import Path
 from PIL import Image
 import random, json
 from modutils import (safe_float, escape_lora_basename, to_lora_key, to_lora_path,
     get_local_model_list, get_private_lora_model_lists, get_valid_lora_name,
-    get_valid_lora_path, get_valid_lora_wt, get_lora_info, CIVITAI_SORT, CIVITAI_PERIOD,
-    normalize_prompt_list, get_civitai_info, search_lora_on_civitai, translate_to_en)
+    get_valid_lora_path, get_valid_lora_wt, get_lora_info, CIVITAI_SORT, CIVITAI_PERIOD, CIVITAI_BASEMODEL,
+    normalize_prompt_list, get_civitai_info, search_lora_on_civitai, translate_to_en, get_t2i_model_info, get_civitai_tag, save_image_history)
+
 
-sd_gen = GuiSD()
 #@spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
            model_name = load_diffusers_format_model[0], lora1 = None, lora1_wt = 1.0, lora2 = None, lora2_wt = 1.0,
@@ -801,7 +685,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     gpu_duration = 59
 
     images: list[tuple[PIL.Image.Image, str | None]] = []
-    info: str = ""
+    info_state = info_images = ""
     progress(0, desc="Preparing...")
 
     if randomize_seed:
@@ -828,7 +712,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     sd_gen.load_new_model(model_name, vae, TASK_MODEL_LIST[0])
     progress(1, desc="Model loaded.")
     progress(0, desc="Starting Inference...")
-    images, info = sd_gen_generate_pipeline(prompt, negative_prompt, 1, num_inference_steps,
+    info_state, images, info_images = sd_gen_generate_pipeline(prompt, negative_prompt, 1, num_inference_steps,
         guidance_scale, True, generator, lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt,
         lora4, lora4_wt, lora5, lora5_wt, sampler,
         height, width, model_name, vae, TASK_MODEL_LIST[0], None, "Canny", 512, 1024,
@@ -1008,14 +892,14 @@ def update_lora_dict(path: str):
 def download_lora(dl_urls: str):
     global loras_url_to_path_dict
     dl_path = ""
-    before = get_local_model_list(directory_loras)
+    before = get_local_model_list(DIRECTORY_LORAS)
     urls = []
     for url in [url.strip() for url in dl_urls.split(',')]:
-        local_path = f"{directory_loras}/{url.split('/')[-1]}"
+        local_path = f"{DIRECTORY_LORAS}/{url.split('/')[-1]}"
         if not Path(local_path).exists():
-            download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
+            download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY)
             urls.append(url)
-    after = get_local_model_list(directory_loras)
+    after = get_local_model_list(DIRECTORY_LORAS)
     new_files = list_sub(after, before)
     i = 0
     for file in new_files: