Upload dc.py

dc.py

@@ -40,27 +40,14 @@ from stablepy import logger
 logger.setLevel(logging.CRITICAL)
 
 from env import (
-    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,
-)
+    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": [
@@ -198,16 +185,8 @@ def process_string(input_string):
         return None
 
 ## 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,
-)
+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 Models**
 download_model = ", ".join(download_model_list)
@@ -335,6 +314,7 @@ class GuiSD:
             vae_model=None,
             type_model_precision=torch.float16,
             retain_task_model_in_cache=False,
+            device="cpu",
         )
 
     def infer_short(self, model, pipe_params, progress=gr.Progress(track_tqdm=True)):
@@ -362,6 +342,7 @@ class GuiSD:
             if incompatible_vae:
                 vae_model = None
 
+        self.model.device = torch.device("cpu")
         
         self.model.load_pipe(
             model_name,
@@ -539,6 +520,7 @@ class GuiSD:
                 params_ip_scale.append(scaleip)
 
         # First load
+        self.model.device = torch.device("cuda:0")
         model_precision = torch.float16
         if not self.model:
             from stablepy import Model_Diffusers
@@ -706,29 +688,18 @@ class GuiSD:
             info_state = info_state + "<br>" + vae_msg
         if msg_lora:
             info_state = info_state + "<br>" + "<br>".join(msg_lora)
-        return self.infer_short(self.model, pipe_params), info_state
+        return self.infer_short(self.model, pipe_params, progress), info_state
 ## END MOD
 
 
 from pathlib import Path
-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,
-    normalize_prompt_list,
-    get_civitai_info,
-    search_lora_on_civitai,
-)
+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,
+    normalize_prompt_list, get_civitai_info, search_lora_on_civitai)
 
 sd_gen = GuiSD()
-@spaces.GPU
+#@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,
            lora3 = None, lora3_wt = 1.0, lora4 = None, lora4_wt = 1.0, lora5 = None, lora5_wt = 1.0,
@@ -757,7 +728,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     lora5 = get_valid_lora_path(lora5)
     progress(1, desc="Preparation completed. Starting inference preparation...")
 
-    sd_gen.load_new_model(model_name, vae, task_model_list[0])
+    sd_gen.load_new_model(model_name, vae, task_model_list[0], progress)
     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,
@@ -767,7 +738,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         False, True, 1, True, False, False, False, False, "./images", False, False, False, True, 1, 0.55,
         False, False, False, True, False, "Use same sampler", False, "", "", 0.35, True, True, False, 4, 4, 32,
         False, "", "", 0.35, True, True, False, 4, 4, 32,
-        True, None, None, "plus_face", "original", 0.7, None, None, "base", "style", 0.7
+        True, None, None, "plus_face", "original", 0.7, None, None, "base", "style", 0.7, progress
     )
 
     progress(1, desc="Inference completed.")
@@ -776,7 +747,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     return output_image
 
 
-@spaces.GPU
+#@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,
             lora3 = None, lora3_wt = 1.0, lora4 = None, lora4_wt = 1.0, lora5 = None, lora5_wt = 1.0,