Upload dc.py

dc.py

@@ -1,52 +1,33 @@
 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
-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
+#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)
@@ -57,63 +38,205 @@ 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.DEBUG)
+logger.setLevel(logging.CRITICAL)
 
 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)
+    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
 
+## 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)
@@ -121,13 +244,99 @@ 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, stream=False):
+    def __init__(self):
         self.model = None
-        self.status_loading = False
-        self.sleep_loading = 4
-        self.last_load = datetime.now()
+    
+        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") #
 
     def infer_short(self, model, pipe_params, progress=gr.Progress(track_tqdm=True)):
         #progress(0, desc="Start inference...")
@@ -141,86 +350,31 @@ class GuiSD:
         return img
 
     def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
-        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}"
+        
+        vae_model = vae_model if vae_model != "None" else None
+        model_type = get_model_type(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(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
+                gr.Warning(msg_inc_vae)
 
+        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,
@@ -325,24 +479,23 @@ class GuiSD:
         mode_ip2,
         scale_ip2,
         pag_scale,
+        #progress=gr.Progress(track_tqdm=True),
     ):
-        info_state = html_template_message("Navigating latent space...")
-        #yield info_state, gr.update(), gr.update()
-
+        #progress(0, desc="Preparing inference...")
+        
         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 = []
@@ -365,9 +518,6 @@ 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'")
 
@@ -439,15 +589,15 @@ class GuiSD:
             "high_threshold": high_threshold,
             "value_threshold": value_threshold,
             "distance_threshold": distance_threshold,
-            "lora_A": lora1 if lora1 != "None" else None,
+            "lora_A": lora1 if lora1 != "None" and lora1 != "" else None,
             "lora_scale_A": lora_scale1,
-            "lora_B": lora2 if lora2 != "None" else None,
+            "lora_B": lora2 if lora2 != "None" and lora2 != "" else None,
             "lora_scale_B": lora_scale2,
-            "lora_C": lora3 if lora3 != "None" else None,
+            "lora_C": lora3 if lora3 != "None" and lora3 != "" else None,
             "lora_scale_C": lora_scale3,
-            "lora_D": lora4 if lora4 != "None" else None,
+            "lora_D": lora4 if lora4 != "None" and lora4 != "" else None,
             "lora_scale_D": lora_scale4,
-            "lora_E": lora5 if lora5 != "None" else None,
+            "lora_E": lora5 if lora5 != "None" and lora5 != "" else None,
             "lora_scale_E": lora_scale5,
 ## BEGIN MOD
             "textual_inversion": get_embed_list(self.model.class_name) if textual_inversion else [],
@@ -497,59 +647,18 @@ class GuiSD:
         }
 
         self.model.device = torch.device("cuda:0")
-        if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5:
+        if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5 and loras_list != [""] * 5:
             self.model.pipe.transformer.to(self.model.device)
             print("transformer to cuda")
 
-        #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}"
+        #progress(1, desc="Inference preparation completed. Starting inference...")
 
-## BEGIN MOD
-                img = save_images(img, metadata)
+        info_state = "" # for yield version
+        return self.infer_short(self.model, pipe_params), info_state
 ## 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)
@@ -569,7 +678,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 # MOD
+        None if item == "None" or item == "" else item
         for item in [args[7], args[9], args[11], args[13], args[15]]
     ]
     lora_status = [None] * 5
@@ -579,7 +688,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 msg_load_lora, gr.update(), gr.update()
+    #    yield None, msg_load_lora
 
     # Load lora in CPU
     if load_lora_cpu:
@@ -605,16 +714,14 @@ 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.\nModel: {sd_gen.model.base_model_id}"
-    if verbose_arg:
+        msg_request = f"Requesting {gpu_duration_arg}s. of GPU time"
         gr.Info(msg_request)
         print(msg_request)
-    #yield msg_request.replace("\n", "<br>"), gr.update(), gr.update()
+
+    # yield from sd_gen.generate_pipeline(*generation_args)
 
     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,
@@ -622,19 +729,31 @@ 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)
 
-    yield msg_task_complete, gr.update(), gr.update()
+def extract_exif_data(image):
+    if image is None: return ""
+
+    try:
+        metadata_keys = ['parameters', 'metadata', 'prompt', 'Comment']
+
+        for key in metadata_keys:
+            if key in image.info:
+                return image.info[key]
 
+        return str(image.info)
 
-@spaces.GPU(duration=15)
+    except Exception as e:
+        return f"Error extracting metadata: {str(e)}"
+
+@spaces.GPU(duration=20)
 def esrgan_upscale(image, upscaler_name, upscaler_size):
     if image is None: return None
 
@@ -656,21 +775,18 @@ 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, 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)
-
+    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)
 
+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,
@@ -685,7 +801,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     gpu_duration = 59
 
     images: list[tuple[PIL.Image.Image, str | None]] = []
-    info_state = info_images = ""
+    info: str = ""
     progress(0, desc="Preparing...")
 
     if randomize_seed:
@@ -712,7 +828,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...")
-    info_state, images, info_images = sd_gen_generate_pipeline(prompt, negative_prompt, 1, num_inference_steps,
+    images, info = 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,
@@ -892,14 +1008,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: