Kohaku-Blueleaf commited on
Commit
c739865
1 Parent(s): d06b16f

fix loading model

Browse files
Files changed (2) hide show
  1. app.py +33 -11
  2. meta.py +8 -0
app.py CHANGED
@@ -1,5 +1,6 @@
1
  import random
2
- from time import time_ns
 
3
 
4
  import torch
5
 
@@ -18,12 +19,17 @@ from meta import (
18
  )
19
 
20
 
21
- sdxl_pipe = load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon", device="cuda")
22
  models.load_model(models.model_list[0])
23
- models.text_model.cuda()
 
24
 
25
  current_dtg_model = models.model_list[0]
26
- current_sdxl_model = "KBlueLeaf/Kohaku-XL-Epsilon"
 
 
 
 
27
 
28
 
29
  @spaces.GPU
@@ -35,14 +41,25 @@ def gen(
35
  addon_prompt: str = "",
36
  seed: int = -1,
37
  ):
38
- global current_dtg_model, current_sdxl_model, sdxl_pipe
39
  if sdxl_model != current_sdxl_model:
40
- sdxl_pipe = load_model(model_id=sdxl_model, device="cuda")
41
- current_sdxl_model = sdxl_model
 
 
 
42
  if dtg_model != current_dtg_model:
43
- models.load_model(dtg_model)
44
- models.text_model.cuda()
45
- current_dtg_model = dtg_model
 
 
 
 
 
 
 
 
46
 
47
  t0 = time_ns()
48
  seed = int(seed)
@@ -54,15 +71,17 @@ def gen(
54
  f"{DEFAULT_STYLE_LIST[style]}, "
55
  f"{MODEL_DEFAULT_QUALITY_LIST[sdxl_model]}, "
56
  )
 
57
  full_prompt = process(
58
  prompt,
59
  aspect_ratio=1.0,
60
  seed=seed,
61
- tag_length="short",
62
  ban_tags=".*alternate.*, character doll, multiple.*, .*cosplay.*, .*name, .*text.*",
63
  format=MODEL_FORMAT_LIST[sdxl_model],
64
  temperature=1.0,
65
  )
 
66
  torch.cuda.empty_cache()
67
 
68
  prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
@@ -82,6 +101,9 @@ def gen(
82
  torch.cuda.empty_cache()
83
  t1 = time_ns()
84
 
 
 
 
85
  return (
86
  result.convert("RGB"),
87
  full_prompt,
 
1
  import random
2
+ from time import time_ns, sleep
3
+ from threading import Lock
4
 
5
  import torch
6
 
 
19
  )
20
 
21
 
22
+ sdxl_pipe = load_model(model_id="KBlueLeaf/Kohaku-XL-Epsilon-rev2", device="cuda")
23
  models.load_model(models.model_list[0])
24
+ models.text_model.cpu()
25
+ torch.cuda.empty_cache()
26
 
27
  current_dtg_model = models.model_list[0]
28
+ current_sdxl_model = "KBlueLeaf/Kohaku-XL-Epsilon-rev2"
29
+
30
+ model_loading_lock = Lock()
31
+ model_running_lock = Lock()
32
+ model_running = 0
33
 
34
 
35
  @spaces.GPU
 
41
  addon_prompt: str = "",
42
  seed: int = -1,
43
  ):
44
+ global current_dtg_model, current_sdxl_model, sdxl_pipe, model_running
45
  if sdxl_model != current_sdxl_model:
46
+ with model_loading_lock:
47
+ while model_running:
48
+ sleep(0.01)
49
+ sdxl_pipe = load_model(model_id=sdxl_model, device="cuda")
50
+ current_sdxl_model = sdxl_model
51
  if dtg_model != current_dtg_model:
52
+ with model_loading_lock:
53
+ while model_running:
54
+ sleep(0.01)
55
+ models.load_model(dtg_model)
56
+ current_dtg_model = dtg_model
57
+
58
+ with model_loading_lock:
59
+ pass
60
+
61
+ with model_running_lock:
62
+ model_running += 1
63
 
64
  t0 = time_ns()
65
  seed = int(seed)
 
71
  f"{DEFAULT_STYLE_LIST[style]}, "
72
  f"{MODEL_DEFAULT_QUALITY_LIST[sdxl_model]}, "
73
  )
74
+ models.text_model.cuda()
75
  full_prompt = process(
76
  prompt,
77
  aspect_ratio=1.0,
78
  seed=seed,
79
+ tag_length="long",
80
  ban_tags=".*alternate.*, character doll, multiple.*, .*cosplay.*, .*name, .*text.*",
81
  format=MODEL_FORMAT_LIST[sdxl_model],
82
  temperature=1.0,
83
  )
84
+ models.text_model.cpu()
85
  torch.cuda.empty_cache()
86
 
87
  prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
 
101
  torch.cuda.empty_cache()
102
  t1 = time_ns()
103
 
104
+ with model_running_lock:
105
+ model_running -= 1
106
+
107
  return (
108
  result.convert("RGB"),
109
  full_prompt,
meta.py CHANGED
@@ -8,11 +8,19 @@ DEFAULT_STYLE_LIST = {
8
  }
9
 
10
  MODEL_DEFAULT_QUALITY_LIST = {
 
11
  "KBlueLeaf/Kohaku-XL-Epsilon": "masterpiece, newest, absurdres, safe",
12
  "cagliostrolab/animagine-xl-3.1": "masterpiece, newest, very aesthetic, absurdres, safe",
13
  }
14
 
15
  MODEL_FORMAT_LIST = {
 
 
 
 
 
 
 
16
  "KBlueLeaf/Kohaku-XL-Epsilon": """<|special|>,
17
  <|characters|>, <|copyrights|>,
18
  <|artist|>,
 
8
  }
9
 
10
  MODEL_DEFAULT_QUALITY_LIST = {
11
+ "KBlueLeaf/Kohaku-XL-Epsilon-rev2": "masterpiece, newest, absurdres",
12
  "KBlueLeaf/Kohaku-XL-Epsilon": "masterpiece, newest, absurdres, safe",
13
  "cagliostrolab/animagine-xl-3.1": "masterpiece, newest, very aesthetic, absurdres, safe",
14
  }
15
 
16
  MODEL_FORMAT_LIST = {
17
+ "KBlueLeaf/Kohaku-XL-Epsilon-rev2": """<|special|>,
18
+ <|characters|>, <|copyrights|>,
19
+ <|artist|>,
20
+
21
+ <|general|>,
22
+
23
+ <|quality|>, <|meta|>, <|rating|>""",
24
  "KBlueLeaf/Kohaku-XL-Epsilon": """<|special|>,
25
  <|characters|>, <|copyrights|>,
26
  <|artist|>,