fix loading model

app.py

@@ -1,5 +1,6 @@
 import random
-from time import time_ns
+from time import time_ns, sleep
+from threading import Lock
 
 import torch
 
@@ -18,12 +19,17 @@ from meta import (
 )
 
 
-sdxl_pipe = load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda")
+sdxl_pipe = load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon-rev2", device="cuda")
 models.load_model(models.model_list[0])
-models.text_model.cuda()
+models.text_model.cpu()
+torch.cuda.empty_cache()
 
 current_dtg_model = models.model_list[0]
-current_sdxl_model = "KBlueLeaf/Kohaku-XL-Epsilon"
+current_sdxl_model = "KBlueLeaf/Kohaku-XL-Epsilon-rev2"
+
+model_loading_lock = Lock()
+model_running_lock = Lock()
+model_running = 0
 
 
 @spaces.GPU
@@ -35,14 +41,25 @@ def gen(
     addon_prompt: str = "",
     seed: int = -1,
 ):
-    global current_dtg_model, current_sdxl_model, sdxl_pipe
+    global current_dtg_model, current_sdxl_model, sdxl_pipe, model_running
     if sdxl_model != current_sdxl_model:
-        sdxl_pipe = load_model(model_id=sdxl_model, device="cuda")
-        current_sdxl_model = sdxl_model
+        with model_loading_lock:
+            while model_running:
+                sleep(0.01)
+            sdxl_pipe = load_model(model_id=sdxl_model, device="cuda")
+            current_sdxl_model = sdxl_model
     if dtg_model != current_dtg_model:
-        models.load_model(dtg_model)
-        models.text_model.cuda()
-        current_dtg_model = dtg_model
+        with model_loading_lock:
+            while model_running:
+                sleep(0.01)
+            models.load_model(dtg_model)
+            current_dtg_model = dtg_model
+
+    with model_loading_lock:
+        pass
+
+    with model_running_lock:
+        model_running += 1
 
     t0 = time_ns()
     seed = int(seed)
@@ -54,15 +71,17 @@ def gen(
         f"{DEFAULT_STYLE_LIST[style]}, "
         f"{MODEL_DEFAULT_QUALITY_LIST[sdxl_model]}, "
     )
+    models.text_model.cuda()
     full_prompt = process(
         prompt,
         aspect_ratio=1.0,
         seed=seed,
-        tag_length="short",
+        tag_length="long",
         ban_tags=".*alternate.*, character doll, multiple.*, .*cosplay.*, .*name, .*text.*",
         format=MODEL_FORMAT_LIST[sdxl_model],
         temperature=1.0,
     )
+    models.text_model.cpu()
     torch.cuda.empty_cache()
 
     prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
@@ -82,6 +101,9 @@ def gen(
     torch.cuda.empty_cache()
     t1 = time_ns()
 
+    with model_running_lock:
+        model_running -= 1
+
     return (
         result.convert("RGB"),
         full_prompt,

meta.py

@@ -8,11 +8,19 @@ DEFAULT_STYLE_LIST = {
 }
 
 MODEL_DEFAULT_QUALITY_LIST = {
+    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": "masterpiece, newest, absurdres",
     "KBlueLeaf/Kohaku-XL-Epsilon": "masterpiece, newest, absurdres, safe",
     "cagliostrolab/animagine-xl-3.1": "masterpiece, newest, very aesthetic, absurdres, safe",
 }
 
 MODEL_FORMAT_LIST = {
+    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": """<|special|>, 
+<|characters|>, <|copyrights|>, 
+<|artist|>, 
+
+<|general|>, 
+
+<|quality|>, <|meta|>, <|rating|>""",
     "KBlueLeaf/Kohaku-XL-Epsilon": """<|special|>, 
 <|characters|>, <|copyrights|>, 
 <|artist|>,