lollms-client 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of lollms-client might be problematic. Click here for more details.
- lollms_client/__init__.py +1 -1
- lollms_client/llm_bindings/azure_openai/__init__.py +2 -2
- lollms_client/llm_bindings/claude/__init__.py +2 -2
- lollms_client/llm_bindings/gemini/__init__.py +2 -2
- lollms_client/llm_bindings/grok/__init__.py +2 -2
- lollms_client/llm_bindings/groq/__init__.py +2 -2
- lollms_client/llm_bindings/hugging_face_inference_api/__init__.py +2 -2
- lollms_client/llm_bindings/litellm/__init__.py +1 -1
- lollms_client/llm_bindings/llamacpp/__init__.py +2 -2
- lollms_client/llm_bindings/lollms/__init__.py +1 -1
- lollms_client/llm_bindings/lollms_webui/__init__.py +1 -1
- lollms_client/llm_bindings/mistral/__init__.py +2 -2
- lollms_client/llm_bindings/novita_ai/__init__.py +2 -2
- lollms_client/llm_bindings/ollama/__init__.py +7 -4
- lollms_client/llm_bindings/open_router/__init__.py +2 -2
- lollms_client/llm_bindings/openai/__init__.py +1 -1
- lollms_client/llm_bindings/openllm/__init__.py +2 -2
- lollms_client/llm_bindings/openwebui/__init__.py +1 -1
- lollms_client/llm_bindings/perplexity/__init__.py +2 -2
- lollms_client/llm_bindings/pythonllamacpp/__init__.py +3 -3
- lollms_client/llm_bindings/tensor_rt/__init__.py +1 -1
- lollms_client/llm_bindings/transformers/__init__.py +4 -4
- lollms_client/llm_bindings/vllm/__init__.py +1 -1
- lollms_client/lollms_core.py +4 -1442
- lollms_client/lollms_llm_binding.py +1 -1
- lollms_client/lollms_tti_binding.py +1 -1
- lollms_client/tti_bindings/diffusers/__init__.py +276 -856
- lollms_client/tti_bindings/diffusers/server/main.py +730 -0
- lollms_client/tti_bindings/gemini/__init__.py +1 -1
- lollms_client/tti_bindings/leonardo_ai/__init__.py +1 -1
- lollms_client/tti_bindings/novita_ai/__init__.py +1 -1
- lollms_client/tti_bindings/stability_ai/__init__.py +1 -1
- lollms_client/tts_bindings/lollms/__init__.py +6 -1
- lollms_client/tts_bindings/piper_tts/__init__.py +1 -1
- lollms_client/tts_bindings/xtts/__init__.py +20 -14
- lollms_client/tts_bindings/xtts/server/main.py +17 -4
- {lollms_client-1.6.2.dist-info → lollms_client-1.6.4.dist-info}/METADATA +2 -2
- {lollms_client-1.6.2.dist-info → lollms_client-1.6.4.dist-info}/RECORD +41 -40
- {lollms_client-1.6.2.dist-info → lollms_client-1.6.4.dist-info}/WHEEL +0 -0
- {lollms_client-1.6.2.dist-info → lollms_client-1.6.4.dist-info}/licenses/LICENSE +0 -0
- {lollms_client-1.6.2.dist-info → lollms_client-1.6.4.dist-info}/top_level.txt +0 -0
|
@@ -1,887 +1,307 @@
|
|
|
1
|
-
# lollms_client/tti_bindings/diffusers/__init__.py
|
|
2
1
|
import os
|
|
3
|
-
import
|
|
4
|
-
from io import BytesIO
|
|
5
|
-
from typing import Optional, List, Dict, Any, Union, Tuple
|
|
6
|
-
from pathlib import Path
|
|
2
|
+
import sys
|
|
7
3
|
import base64
|
|
8
|
-
import pipmaster as pm
|
|
9
|
-
import threading
|
|
10
|
-
import queue
|
|
11
|
-
from concurrent.futures import Future
|
|
12
|
-
import time
|
|
13
|
-
import hashlib
|
|
14
4
|
import requests
|
|
15
|
-
|
|
5
|
+
import subprocess
|
|
6
|
+
import time
|
|
16
7
|
import json
|
|
17
|
-
import
|
|
18
|
-
from
|
|
19
|
-
from
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
pm.ensure_packages(["pillow","transformers","safetensors","requests","tqdm"])
|
|
23
|
-
pm.ensure_packages([
|
|
24
|
-
{
|
|
25
|
-
"name": "diffusers",
|
|
26
|
-
"vcs": "git+https://github.com/huggingface/diffusers.git",
|
|
27
|
-
"condition": ">=0.35.1"
|
|
28
|
-
}
|
|
29
|
-
])
|
|
8
|
+
from io import BytesIO
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import Optional, List, Dict, Any, Union
|
|
11
|
+
|
|
12
|
+
# Ensure pipmaster is available.
|
|
30
13
|
try:
|
|
31
|
-
import
|
|
32
|
-
from diffusers import (
|
|
33
|
-
AutoPipelineForText2Image,
|
|
34
|
-
AutoPipelineForImage2Image,
|
|
35
|
-
AutoPipelineForInpainting,
|
|
36
|
-
DiffusionPipeline,
|
|
37
|
-
StableDiffusionPipeline,
|
|
38
|
-
QwenImageEditPipeline,
|
|
39
|
-
QwenImageEditPlusPipeline
|
|
40
|
-
)
|
|
41
|
-
from diffusers.utils import load_image
|
|
42
|
-
from PIL import Image
|
|
43
|
-
DIFFUSERS_AVAILABLE = True
|
|
14
|
+
import pipmaster as pm
|
|
44
15
|
except ImportError:
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
"
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
"
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
"
|
|
90
|
-
"url": "https://civitai.com/api/download/models/72396?type=Model&format=SafeTensor&size=full&fp=fp16",
|
|
91
|
-
"filename": "lyriel_v16.safetensors",
|
|
92
|
-
"description": "Fantasy/stylized SD1.5.",
|
|
93
|
-
"owned_by": "civitai"
|
|
94
|
-
},
|
|
95
|
-
"ui_icons": {
|
|
96
|
-
"display_name": "UI Icons",
|
|
97
|
-
"url": "https://civitai.com/api/download/models/367044?type=Model&format=SafeTensor&size=full&fp=fp16",
|
|
98
|
-
"filename": "uiIcons_v10.safetensors",
|
|
99
|
-
"description": "A model for generating UI icons.",
|
|
100
|
-
"owned_by": "civitai"
|
|
101
|
-
},
|
|
102
|
-
"meinamix": {
|
|
103
|
-
"display_name": "MeinaMix",
|
|
104
|
-
"url": "https://civitai.com/api/download/models/948574?type=Model&format=SafeTensor&size=pruned&fp=fp16",
|
|
105
|
-
"filename": "meinamix_meinaV11.safetensors",
|
|
106
|
-
"description": "Anime/illustration SD1.5.",
|
|
107
|
-
"owned_by": "civitai"
|
|
108
|
-
},
|
|
109
|
-
"rpg-v5": {
|
|
110
|
-
"display_name": "RPG v5",
|
|
111
|
-
"url": "https://civitai.com/api/download/models/124626?type=Model&format=SafeTensor&size=pruned&fp=fp16",
|
|
112
|
-
"filename": "rpg_v5.safetensors",
|
|
113
|
-
"description": "RPG assets SD1.5.",
|
|
114
|
-
"owned_by": "civitai"
|
|
115
|
-
},
|
|
116
|
-
"pixel-art-xl": {
|
|
117
|
-
"display_name": "Pixel Art XL",
|
|
118
|
-
"url": "https://civitai.com/api/download/models/135931?type=Model&format=SafeTensor",
|
|
119
|
-
"filename": "pixelartxl_v11.safetensors",
|
|
120
|
-
"description": "Pixel art SDXL.",
|
|
121
|
-
"owned_by": "civitai"
|
|
122
|
-
},
|
|
123
|
-
"lowpoly-world": {
|
|
124
|
-
"display_name": "Lowpoly World",
|
|
125
|
-
"url": "https://civitai.com/api/download/models/146502?type=Model&format=SafeTensor",
|
|
126
|
-
"filename": "LowpolySDXL.safetensors",
|
|
127
|
-
"description": "Lowpoly style SD1.5.",
|
|
128
|
-
"owned_by": "civitai"
|
|
129
|
-
},
|
|
130
|
-
"toonyou": {
|
|
131
|
-
"display_name": "ToonYou",
|
|
132
|
-
"url": "https://civitai.com/api/download/models/125771?type=Model&format=SafeTensor&size=pruned&fp=fp16",
|
|
133
|
-
"filename": "toonyou_beta6.safetensors",
|
|
134
|
-
"description": "Cartoon/Disney SD1.5.",
|
|
135
|
-
"owned_by": "civitai"
|
|
136
|
-
},
|
|
137
|
-
"papercut": {
|
|
138
|
-
"display_name": "Papercut",
|
|
139
|
-
"url": "https://civitai.com/api/download/models/133503?type=Model&format=SafeTensor",
|
|
140
|
-
"filename": "papercut.safetensors",
|
|
141
|
-
"description": "Paper cutout SD1.5.",
|
|
142
|
-
"owned_by": "civitai"
|
|
143
|
-
},
|
|
144
|
-
"fantassifiedIcons": {
|
|
145
|
-
"display_name": "Fantassified Icons",
|
|
146
|
-
"url": "https://civitai.com/api/download/models/67584?type=Model&format=SafeTensor&size=pruned&fp=fp16",
|
|
147
|
-
"filename": "fantassifiedIcons_fantassifiedIconsV20.safetensors",
|
|
148
|
-
"description": "Flat, modern Icons.",
|
|
149
|
-
"owned_by": "civitai"
|
|
150
|
-
},
|
|
151
|
-
"game_icon_institute": {
|
|
152
|
-
"display_name": "Game icon institute",
|
|
153
|
-
"url": "https://civitai.com/api/download/models/158776?type=Model&format=SafeTensor&size=full&fp=fp16",
|
|
154
|
-
"filename": "gameIconInstituteV10_v10.safetensors",
|
|
155
|
-
"description": "Flat, modern game Icons.",
|
|
156
|
-
"owned_by": "civitai"
|
|
157
|
-
},
|
|
158
|
-
"M4RV3LS_DUNGEONS": {
|
|
159
|
-
"display_name": "M4RV3LS & DUNGEONS",
|
|
160
|
-
"url": "https://civitai.com/api/download/models/139417?type=Model&format=SafeTensor&size=pruned&fp=fp16",
|
|
161
|
-
"filename": "M4RV3LSDUNGEONSNEWV40COMICS_mD40.safetensors",
|
|
162
|
-
"description": "comics.",
|
|
163
|
-
"owned_by": "civitai"
|
|
164
|
-
},
|
|
165
|
-
}
|
|
166
|
-
|
|
167
|
-
TORCH_DTYPE_MAP_STR_TO_OBJ = {
|
|
168
|
-
"float16": getattr(torch, 'float16', 'float16'),
|
|
169
|
-
"bfloat16": getattr(torch, 'bfloat16', 'bfloat16'),
|
|
170
|
-
"float32": getattr(torch, 'float32', 'float32'),
|
|
171
|
-
"auto": "auto"
|
|
172
|
-
}
|
|
173
|
-
TORCH_DTYPE_MAP_OBJ_TO_STR = {v: k for k, v in TORCH_DTYPE_MAP_STR_TO_OBJ.items()}
|
|
174
|
-
if torch:
|
|
175
|
-
TORCH_DTYPE_MAP_OBJ_TO_STR[None] = "None"
|
|
176
|
-
|
|
177
|
-
SCHEDULER_MAPPING = {
|
|
178
|
-
"default": None,
|
|
179
|
-
"ddim": "DDIMScheduler",
|
|
180
|
-
"ddpm": "DDPMScheduler",
|
|
181
|
-
"deis_multistep": "DEISMultistepScheduler",
|
|
182
|
-
"dpm_multistep": "DPMSolverMultistepScheduler",
|
|
183
|
-
"dpm_multistep_karras": "DPMSolverMultistepScheduler",
|
|
184
|
-
"dpm_single": "DPMSolverSinglestepScheduler",
|
|
185
|
-
"dpm_adaptive": "DPMSolverPlusPlusScheduler",
|
|
186
|
-
"dpm++_2m": "DPMSolverMultistepScheduler",
|
|
187
|
-
"dpm++_2m_karras": "DPMSolverMultistepScheduler",
|
|
188
|
-
"dpm++_2s_ancestral": "DPMSolverAncestralDiscreteScheduler",
|
|
189
|
-
"dpm++_2s_ancestral_karras": "DPMSolverAncestralDiscreteScheduler",
|
|
190
|
-
"dpm++_sde": "DPMSolverSDEScheduler",
|
|
191
|
-
"dpm++_sde_karras": "DPMSolverSDEScheduler",
|
|
192
|
-
"euler_ancestral_discrete": "EulerAncestralDiscreteScheduler",
|
|
193
|
-
"euler_discrete": "EulerDiscreteScheduler",
|
|
194
|
-
"heun_discrete": "HeunDiscreteScheduler",
|
|
195
|
-
"heun_karras": "HeunDiscreteScheduler",
|
|
196
|
-
"lms_discrete": "LMSDiscreteScheduler",
|
|
197
|
-
"lms_karras": "LMSDiscreteScheduler",
|
|
198
|
-
"pndm": "PNDMScheduler",
|
|
199
|
-
"unipc_multistep": "UniPCMultistepScheduler",
|
|
200
|
-
"dpm++_2m_sde": "DPMSolverMultistepScheduler",
|
|
201
|
-
"dpm++_2m_sde_karras": "DPMSolverMultistepScheduler",
|
|
202
|
-
"dpm2": "KDPM2DiscreteScheduler",
|
|
203
|
-
"dpm2_karras": "KDPM2DiscreteScheduler",
|
|
204
|
-
"dpm2_a": "KDPM2AncestralDiscreteScheduler",
|
|
205
|
-
"dpm2_a_karras": "KDPM2AncestralDiscreteScheduler",
|
|
206
|
-
"euler": "EulerDiscreteScheduler",
|
|
207
|
-
"euler_a": "EulerAncestralDiscreteScheduler",
|
|
208
|
-
"heun": "HeunDiscreteScheduler",
|
|
209
|
-
"lms": "LMSDiscreteScheduler"
|
|
210
|
-
}
|
|
211
|
-
SCHEDULER_USES_KARRAS_SIGMAS = [
|
|
212
|
-
"dpm_multistep_karras","dpm++_2m_karras","dpm++_2s_ancestral_karras",
|
|
213
|
-
"dpm++_sde_karras","heun_karras","lms_karras",
|
|
214
|
-
"dpm++_2m_sde_karras","dpm2_karras","dpm2_a_karras"
|
|
215
|
-
]
|
|
216
|
-
|
|
217
|
-
class ModelManager:
|
|
218
|
-
def __init__(self, config: Dict[str, Any], models_path: Path):
|
|
219
|
-
self.config = config
|
|
220
|
-
self.models_path = models_path
|
|
221
|
-
self.pipeline: Optional[DiffusionPipeline] = None
|
|
222
|
-
self.current_task: Optional[str] = None
|
|
223
|
-
self.ref_count = 0
|
|
224
|
-
self.lock = threading.Lock()
|
|
225
|
-
self.queue = queue.Queue()
|
|
226
|
-
self.is_loaded = False
|
|
227
|
-
self.last_used_time = time.time()
|
|
228
|
-
self._stop_event = threading.Event()
|
|
229
|
-
self.worker_thread = threading.Thread(target=self._generation_worker, daemon=True)
|
|
230
|
-
self.worker_thread.start()
|
|
231
|
-
self._stop_monitor_event = threading.Event()
|
|
232
|
-
self._unload_monitor_thread = None
|
|
233
|
-
self._start_unload_monitor()
|
|
234
|
-
|
|
235
|
-
def acquire(self):
|
|
236
|
-
with self.lock:
|
|
237
|
-
self.ref_count += 1
|
|
238
|
-
return self
|
|
239
|
-
|
|
240
|
-
def release(self):
|
|
241
|
-
with self.lock:
|
|
242
|
-
self.ref_count -= 1
|
|
243
|
-
return self.ref_count
|
|
244
|
-
|
|
245
|
-
def stop(self):
|
|
246
|
-
self._stop_event.set()
|
|
247
|
-
if self._unload_monitor_thread:
|
|
248
|
-
self._stop_monitor_event.set()
|
|
249
|
-
self._unload_monitor_thread.join(timeout=2)
|
|
250
|
-
self.queue.put(None)
|
|
251
|
-
self.worker_thread.join(timeout=5)
|
|
252
|
-
|
|
253
|
-
def _start_unload_monitor(self):
|
|
254
|
-
unload_after = self.config.get("unload_inactive_model_after", 0)
|
|
255
|
-
if unload_after > 0 and self._unload_monitor_thread is None:
|
|
256
|
-
self._stop_monitor_event.clear()
|
|
257
|
-
self._unload_monitor_thread = threading.Thread(target=self._unload_monitor, daemon=True)
|
|
258
|
-
self._unload_monitor_thread.start()
|
|
259
|
-
|
|
260
|
-
def _unload_monitor(self):
|
|
261
|
-
unload_after = self.config.get("unload_inactive_model_after", 0)
|
|
262
|
-
if unload_after <= 0:
|
|
263
|
-
return
|
|
264
|
-
ASCIIColors.info(f"Starting inactivity monitor for '{self.config['model_name']}' (timeout: {unload_after}s).")
|
|
265
|
-
while not self._stop_monitor_event.wait(timeout=5.0):
|
|
266
|
-
with self.lock:
|
|
267
|
-
if not self.is_loaded:
|
|
268
|
-
continue
|
|
269
|
-
if time.time() - self.last_used_time > unload_after:
|
|
270
|
-
ASCIIColors.info(f"Model '{self.config['model_name']}' has been inactive. Unloading.")
|
|
271
|
-
self._unload_pipeline()
|
|
272
|
-
|
|
273
|
-
def _resolve_model_path(self, model_name: str) -> Union[str, Path]:
|
|
274
|
-
path_obj = Path(model_name)
|
|
275
|
-
if path_obj.is_absolute() and path_obj.exists():
|
|
276
|
-
return model_name
|
|
277
|
-
if model_name in CIVITAI_MODELS:
|
|
278
|
-
filename = CIVITAI_MODELS[model_name]["filename"]
|
|
279
|
-
local_path = self.models_path / filename
|
|
280
|
-
if not local_path.exists():
|
|
281
|
-
self._download_civitai_model(model_name)
|
|
282
|
-
return local_path
|
|
283
|
-
local_path = self.models_path / model_name
|
|
284
|
-
if local_path.exists():
|
|
285
|
-
return local_path
|
|
286
|
-
return model_name
|
|
287
|
-
|
|
288
|
-
def _download_civitai_model(self, model_key: str):
|
|
289
|
-
model_info = CIVITAI_MODELS[model_key]
|
|
290
|
-
url = model_info["url"]
|
|
291
|
-
filename = model_info["filename"]
|
|
292
|
-
dest_path = self.models_path / filename
|
|
293
|
-
temp_path = dest_path.with_suffix(".temp")
|
|
294
|
-
ASCIIColors.cyan(f"Downloading '{filename}' from Civitai...")
|
|
16
|
+
print("FATAL: pipmaster is not installed. Please install it using: pip install pipmaster")
|
|
17
|
+
sys.exit(1)
|
|
18
|
+
|
|
19
|
+
# Ensure filelock is available for process-safe server startup.
|
|
20
|
+
try:
|
|
21
|
+
from filelock import FileLock, Timeout
|
|
22
|
+
except ImportError:
|
|
23
|
+
print("FATAL: The 'filelock' library is required. Please install it by running: pip install filelock")
|
|
24
|
+
sys.exit(1)
|
|
25
|
+
|
|
26
|
+
from lollms_client.lollms_tti_binding import LollmsTTIBinding
|
|
27
|
+
from ascii_colors import ASCIIColors
|
|
28
|
+
|
|
29
|
+
BindingName = "DiffusersBinding"
|
|
30
|
+
|
|
31
|
+
class DiffusersBinding(LollmsTTIBinding):
|
|
32
|
+
"""
|
|
33
|
+
Client binding for a dedicated, managed Diffusers server.
|
|
34
|
+
This architecture prevents multiple models from being loaded into memory
|
|
35
|
+
in a multi-worker environment, solving OOM errors.
|
|
36
|
+
"""
|
|
37
|
+
def __init__(self,
|
|
38
|
+
**kwargs):
|
|
39
|
+
|
|
40
|
+
super().__init__(binding_name=BindingName)
|
|
41
|
+
|
|
42
|
+
# Prioritize 'model_name' but accept 'model' as an alias from config files.
|
|
43
|
+
if 'model' in kwargs and 'model_name' not in kwargs:
|
|
44
|
+
kwargs['model_name'] = kwargs.pop('model')
|
|
45
|
+
|
|
46
|
+
self.config = kwargs
|
|
47
|
+
self.host = kwargs.get("host", "localhost")
|
|
48
|
+
self.port = kwargs.get("port", 9630)
|
|
49
|
+
self.auto_start_server = kwargs.get("auto_start_server", True)
|
|
50
|
+
self.server_process = None
|
|
51
|
+
self.base_url = f"http://{self.host}:{self.port}"
|
|
52
|
+
self.binding_root = Path(__file__).parent
|
|
53
|
+
self.server_dir = self.binding_root / "server"
|
|
54
|
+
self.venv_dir = Path("./venv/tti_diuffusers_venv")
|
|
55
|
+
self.models_path = Path(kwargs.get("models_path", "./diffusers_models")).resolve()
|
|
56
|
+
if self.auto_start_server:
|
|
57
|
+
self.ensure_server_is_running()
|
|
58
|
+
|
|
59
|
+
def is_server_running(self) -> bool:
|
|
60
|
+
"""Checks if the server is already running and responsive."""
|
|
295
61
|
try:
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
scheduler_name_key = self.config["scheduler_name"].lower()
|
|
314
|
-
if scheduler_name_key == "default":
|
|
315
|
-
return
|
|
316
|
-
scheduler_class_name = SCHEDULER_MAPPING.get(scheduler_name_key)
|
|
317
|
-
if scheduler_class_name:
|
|
318
|
-
try:
|
|
319
|
-
SchedulerClass = getattr(importlib.import_module("diffusers.schedulers"), scheduler_class_name)
|
|
320
|
-
scheduler_config = self.pipeline.scheduler.config
|
|
321
|
-
scheduler_config["use_karras_sigmas"] = scheduler_name_key in SCHEDULER_USES_KARRAS_SIGMAS
|
|
322
|
-
self.pipeline.scheduler = SchedulerClass.from_config(scheduler_config)
|
|
323
|
-
ASCIIColors.info(f"Switched scheduler to {scheduler_class_name}")
|
|
324
|
-
except Exception as e:
|
|
325
|
-
ASCIIColors.warning(f"Could not switch scheduler to {scheduler_name_key}: {e}. Using current default.")
|
|
326
|
-
|
|
327
|
-
def _load_pipeline_for_task(self, task: str):
|
|
328
|
-
if self.pipeline and self.current_task == task:
|
|
329
|
-
return
|
|
330
|
-
if self.pipeline:
|
|
331
|
-
self._unload_pipeline()
|
|
332
|
-
model_name = self.config.get("model_name", "")
|
|
333
|
-
if not model_name:
|
|
334
|
-
raise ValueError("Model name cannot be empty for loading.")
|
|
335
|
-
ASCIIColors.info(f"Loading Diffusers model: {model_name} for task: {task}")
|
|
336
|
-
model_path = self._resolve_model_path(model_name)
|
|
337
|
-
torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower())
|
|
62
|
+
response = requests.get(f"{self.base_url}/status", timeout=2)
|
|
63
|
+
if response.status_code == 200 and response.json().get("status") == "running":
|
|
64
|
+
return True
|
|
65
|
+
except requests.exceptions.RequestException:
|
|
66
|
+
return False
|
|
67
|
+
return False
|
|
68
|
+
|
|
69
|
+
def ensure_server_is_running(self):
|
|
70
|
+
"""
|
|
71
|
+
Ensures the Diffusers server is running. If not, it attempts to start it
|
|
72
|
+
in a process-safe manner using a file lock.
|
|
73
|
+
"""
|
|
74
|
+
self.server_dir.mkdir(exist_ok=True)
|
|
75
|
+
lock_path = self.server_dir / "diffusers_server.lock"
|
|
76
|
+
lock = FileLock(lock_path, timeout=60) # Increased timeout for long installs
|
|
77
|
+
|
|
78
|
+
ASCIIColors.info("Attempting to start or connect to the Diffusers server...")
|
|
338
79
|
try:
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
|
|
354
|
-
|
|
355
|
-
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
if not self.config["safety_checker_on"]:
|
|
362
|
-
common_args["safety_checker"] = None
|
|
363
|
-
if self.config.get("hf_cache_path"):
|
|
364
|
-
common_args["cache_dir"] = str(self.config["hf_cache_path"])
|
|
365
|
-
|
|
366
|
-
if "Qwen-Image-Edit-2509" in str(model_path):
|
|
367
|
-
self.pipeline = QwenImageEditPlusPipeline.from_pretrained(model_path, **common_args)
|
|
368
|
-
elif "Qwen-Image-Edit" in str(model_path):
|
|
369
|
-
self.pipeline = QwenImageEditPipeline.from_pretrained(model_path, **common_args)
|
|
370
|
-
elif task == "text2image":
|
|
371
|
-
self.pipeline = AutoPipelineForText2Image.from_pretrained(model_path, **common_args)
|
|
372
|
-
elif task == "image2image":
|
|
373
|
-
self.pipeline = AutoPipelineForImage2Image.from_pretrained(model_path, **common_args)
|
|
374
|
-
elif task == "inpainting":
|
|
375
|
-
self.pipeline = AutoPipelineForInpainting.from_pretrained(model_path, **common_args)
|
|
376
|
-
except Exception as e:
|
|
377
|
-
error_str = str(e).lower()
|
|
378
|
-
if "401" in error_str or "gated" in error_str or "authorization" in error_str:
|
|
379
|
-
msg = (
|
|
380
|
-
f"AUTHENTICATION FAILED for model '{model_name}'. "
|
|
381
|
-
"Please ensure you accepted the model license and provided a valid HF token."
|
|
382
|
-
)
|
|
383
|
-
raise RuntimeError(msg) from e
|
|
384
|
-
raise e
|
|
385
|
-
self._set_scheduler()
|
|
386
|
-
self.pipeline.to(self.config["device"])
|
|
387
|
-
if self.config["enable_xformers"]:
|
|
80
|
+
with lock:
|
|
81
|
+
if not self.is_server_running():
|
|
82
|
+
ASCIIColors.yellow("Lock acquired. Starting dedicated Diffusers server...")
|
|
83
|
+
self.start_server()
|
|
84
|
+
else:
|
|
85
|
+
ASCIIColors.green("Server was started by another process. Connected successfully.")
|
|
86
|
+
except Timeout:
|
|
87
|
+
ASCIIColors.yellow("Could not acquire lock. Another process is likely starting the server. Waiting for it to become available...")
|
|
88
|
+
|
|
89
|
+
self._wait_for_server()
|
|
90
|
+
|
|
91
|
+
def install_server_dependencies(self):
|
|
92
|
+
"""
|
|
93
|
+
Installs the server's dependencies into a dedicated virtual environment
|
|
94
|
+
using pipmaster, which handles complex packages like PyTorch.
|
|
95
|
+
"""
|
|
96
|
+
ASCIIColors.info(f"Setting up virtual environment in: {self.venv_dir}")
|
|
97
|
+
pm_v = pm.PackageManager(venv_path=str(self.venv_dir))
|
|
98
|
+
|
|
99
|
+
# --- PyTorch Installation ---
|
|
100
|
+
torch_index_url = None
|
|
101
|
+
if sys.platform == "win32":
|
|
388
102
|
try:
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
ASCIIColors.
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
395
|
-
|
|
396
|
-
self.is_loaded = True
|
|
397
|
-
self.current_task = task
|
|
398
|
-
self.last_used_time = time.time()
|
|
399
|
-
ASCIIColors.green(f"Model '{model_name}' loaded successfully on '{self.config['device']}' for task '{task}'.")
|
|
400
|
-
|
|
401
|
-
def _unload_pipeline(self):
|
|
402
|
-
if self.pipeline:
|
|
403
|
-
model_name = self.config.get('model_name', 'Unknown')
|
|
404
|
-
del self.pipeline
|
|
405
|
-
self.pipeline = None
|
|
406
|
-
if torch and torch.cuda.is_available():
|
|
407
|
-
torch.cuda.empty_cache()
|
|
408
|
-
self.is_loaded = False
|
|
409
|
-
self.current_task = None
|
|
410
|
-
ASCIIColors.info(f"Model '{model_name}' unloaded and VRAM cleared.")
|
|
411
|
-
|
|
412
|
-
def _generation_worker(self):
|
|
413
|
-
while not self._stop_event.is_set():
|
|
414
|
-
try:
|
|
415
|
-
job = self.queue.get(timeout=1)
|
|
416
|
-
if job is None:
|
|
417
|
-
break
|
|
418
|
-
future, task, pipeline_args = job
|
|
419
|
-
try:
|
|
420
|
-
with self.lock:
|
|
421
|
-
self.last_used_time = time.time()
|
|
422
|
-
if not self.is_loaded or self.current_task != task:
|
|
423
|
-
self._load_pipeline_for_task(task)
|
|
424
|
-
with torch.no_grad():
|
|
425
|
-
output = self.pipeline(**pipeline_args)
|
|
426
|
-
pil = output.images[0]
|
|
427
|
-
buf = BytesIO()
|
|
428
|
-
pil.save(buf, format="PNG")
|
|
429
|
-
future.set_result(buf.getvalue())
|
|
430
|
-
except Exception as e:
|
|
431
|
-
trace_exception(e)
|
|
432
|
-
future.set_exception(e)
|
|
433
|
-
finally:
|
|
434
|
-
self.queue.task_done()
|
|
435
|
-
except queue.Empty:
|
|
436
|
-
continue
|
|
437
|
-
|
|
438
|
-
class PipelineRegistry:
|
|
439
|
-
_instance = None
|
|
440
|
-
_lock = threading.Lock()
|
|
441
|
-
def __new__(cls, *args, **kwargs):
|
|
442
|
-
with cls._lock:
|
|
443
|
-
if cls._instance is None:
|
|
444
|
-
cls._instance = super().__new__(cls)
|
|
445
|
-
cls._instance._managers = {}
|
|
446
|
-
cls._instance._registry_lock = threading.Lock()
|
|
447
|
-
return cls._instance
|
|
448
|
-
@staticmethod
|
|
449
|
-
def _get_critical_keys():
|
|
450
|
-
return [
|
|
451
|
-
"model_name","device","torch_dtype_str","use_safetensors",
|
|
452
|
-
"safety_checker_on","hf_variant","enable_cpu_offload",
|
|
453
|
-
"enable_sequential_cpu_offload","enable_xformers",
|
|
454
|
-
"local_files_only","hf_cache_path","unload_inactive_model_after"
|
|
455
|
-
]
|
|
456
|
-
def _get_config_key(self, config: Dict[str, Any]) -> str:
|
|
457
|
-
key_data = tuple(sorted((k, config.get(k)) for k in self._get_critical_keys()))
|
|
458
|
-
return hashlib.sha256(str(key_data).encode('utf-8')).hexdigest()
|
|
459
|
-
def get_manager(self, config: Dict[str, Any], models_path: Path) -> ModelManager:
|
|
460
|
-
key = self._get_config_key(config)
|
|
461
|
-
with self._registry_lock:
|
|
462
|
-
if key not in self._managers:
|
|
463
|
-
self._managers[key] = ModelManager(config.copy(), models_path)
|
|
464
|
-
return self._managers[key].acquire()
|
|
465
|
-
def release_manager(self, config: Dict[str, Any]):
|
|
466
|
-
key = self._get_config_key(config)
|
|
467
|
-
with self._registry_lock:
|
|
468
|
-
if key in self._managers:
|
|
469
|
-
manager = self._managers[key]
|
|
470
|
-
ref_count = manager.release()
|
|
471
|
-
if ref_count == 0:
|
|
472
|
-
ASCIIColors.info(f"Reference count for model '{config.get('model_name')}' is zero. Cleaning up manager.")
|
|
473
|
-
manager.stop()
|
|
474
|
-
with manager.lock:
|
|
475
|
-
manager._unload_pipeline()
|
|
476
|
-
del self._managers[key]
|
|
477
|
-
def get_active_managers(self) -> List[ModelManager]:
|
|
478
|
-
with self._registry_lock:
|
|
479
|
-
return [m for m in self._managers.values() if m.is_loaded]
|
|
480
|
-
|
|
481
|
-
class DiffusersTTIBinding_Impl(LollmsTTIBinding):
|
|
482
|
-
DEFAULT_CONFIG = {
|
|
483
|
-
"model_name": "",
|
|
484
|
-
"device": "auto",
|
|
485
|
-
"torch_dtype_str": "auto",
|
|
486
|
-
"use_safetensors": True,
|
|
487
|
-
"scheduler_name": "default",
|
|
488
|
-
"safety_checker_on": True,
|
|
489
|
-
"num_inference_steps": 25,
|
|
490
|
-
"guidance_scale": 7.0,
|
|
491
|
-
"width": 512,
|
|
492
|
-
"height": 512,
|
|
493
|
-
"seed": -1,
|
|
494
|
-
"enable_cpu_offload": False,
|
|
495
|
-
"enable_sequential_cpu_offload": False,
|
|
496
|
-
"enable_xformers": False,
|
|
497
|
-
"hf_variant": None,
|
|
498
|
-
"hf_token": None,
|
|
499
|
-
"hf_cache_path": None,
|
|
500
|
-
"local_files_only": False,
|
|
501
|
-
"unload_inactive_model_after": 0
|
|
502
|
-
}
|
|
503
|
-
HF_DEFAULT_MODELS = [
|
|
504
|
-
{"family": "SDXL", "model_name": "stabilityai/stable-diffusion-xl-base-1.0", "display_name": "SDXL Base 1.0", "desc": "Text2Image 1024 native."},
|
|
505
|
-
{"family": "SDXL", "model_name": "stabilityai/stable-diffusion-xl-refiner-1.0", "display_name": "SDXL Refiner 1.0", "desc": "Refiner for SDXL."},
|
|
506
|
-
{"family": "SD 1.x", "model_name": "runwayml/stable-diffusion-v1-5", "display_name": "Stable Diffusion 1.5", "desc": "Classic SD1.5."},
|
|
507
|
-
{"family": "SD 2.x", "model_name": "stabilityai/stable-diffusion-2-1", "display_name": "Stable Diffusion 2.1", "desc": "SD2.1 base."},
|
|
508
|
-
{"family": "SD3", "model_name": "stabilityai/stable-diffusion-3-medium-diffusers", "display_name": "Stable Diffusion 3 Medium", "desc": "SD3 medium."},
|
|
509
|
-
{"family": "Qwen", "model_name": "Qwen/Qwen-Image", "display_name": "Qwen Image", "desc": "Dedicated image generation."},
|
|
510
|
-
{"family": "Specialized", "model_name": "playgroundai/playground-v2.5-1024px-aesthetic", "display_name": "Playground v2.5", "desc": "High aesthetic 1024."},
|
|
511
|
-
{"family": "Editors", "model_name": "Qwen/Qwen-Image-Edit", "display_name": "Qwen Image Edit", "desc": "Dedicated image editing."},
|
|
512
|
-
{"family": "Editors", "model_name": "Qwen/Qwen-Image-Edit-2509", "display_name": "Qwen Image Edit Plus (Multi-Image)", "desc": "Advanced multi-image editing, fusion, and pose transfer."}
|
|
513
|
-
]
|
|
514
|
-
|
|
515
|
-
def __init__(self, **kwargs):
|
|
516
|
-
super().__init__(binding_name=BindingName)
|
|
517
|
-
self.manager: Optional[ModelManager] = None
|
|
518
|
-
if not DIFFUSERS_AVAILABLE:
|
|
519
|
-
raise ImportError("Diffusers not available. Please install required packages.")
|
|
520
|
-
self.config = self.DEFAULT_CONFIG.copy()
|
|
521
|
-
self.config.update(kwargs)
|
|
522
|
-
self.model_name = self.config.get("model_name", "")
|
|
103
|
+
# Use nvidia-smi to detect CUDA
|
|
104
|
+
result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=True)
|
|
105
|
+
ASCIIColors.green("NVIDIA GPU detected. Installing CUDA-enabled PyTorch.")
|
|
106
|
+
# Using a common and stable CUDA version. Adjust if needed.
|
|
107
|
+
torch_index_url = "https://download.pytorch.org/whl/cu121"
|
|
108
|
+
except (FileNotFoundError, subprocess.CalledProcessError):
|
|
109
|
+
ASCIIColors.yellow("`nvidia-smi` not found or failed. Installing standard PyTorch. If you have an NVIDIA GPU, please ensure drivers are installed and in PATH.")
|
|
523
110
|
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
self.models_path.mkdir(parents=True, exist_ok=True)
|
|
527
|
-
self.registry = PipelineRegistry()
|
|
528
|
-
self._resolve_device_and_dtype()
|
|
529
|
-
if self.model_name:
|
|
530
|
-
self._acquire_manager()
|
|
111
|
+
# Base packages including torch. pm.ensure_packages handles verbose output.
|
|
112
|
+
pm_v.ensure_packages(["torch", "torchvision"], index_url=torch_index_url)
|
|
531
113
|
|
|
532
|
-
def ps(self) -> List[dict]:
|
|
533
|
-
if not self.registry:
|
|
534
|
-
return []
|
|
535
|
-
try:
|
|
536
|
-
active = self.registry.get_active_managers()
|
|
537
|
-
out = []
|
|
538
|
-
for m in active:
|
|
539
|
-
with m.lock:
|
|
540
|
-
cfg = m.config
|
|
541
|
-
pipe = m.pipeline
|
|
542
|
-
vram_usage_bytes = 0
|
|
543
|
-
if torch.cuda.is_available() and cfg.get("device") == "cuda" and pipe:
|
|
544
|
-
for comp in pipe.components.values():
|
|
545
|
-
if hasattr(comp, 'parameters'):
|
|
546
|
-
mem_params = sum(p.nelement() * p.element_size() for p in comp.parameters())
|
|
547
|
-
mem_bufs = sum(b.nelement() * b.element_size() for b in comp.buffers())
|
|
548
|
-
vram_usage_bytes += (mem_params + mem_bufs)
|
|
549
|
-
out.append({
|
|
550
|
-
"model_name": cfg.get("model_name"),
|
|
551
|
-
"vram_size": vram_usage_bytes,
|
|
552
|
-
"device": cfg.get("device"),
|
|
553
|
-
"torch_dtype": str(pipe.dtype) if pipe else cfg.get("torch_dtype_str"),
|
|
554
|
-
"pipeline_type": pipe.__class__.__name__ if pipe else "N/A",
|
|
555
|
-
"scheduler_class": pipe.scheduler.__class__.__name__ if pipe and hasattr(pipe, 'scheduler') else "N/A",
|
|
556
|
-
"status": "Active" if m.is_loaded else "Idle",
|
|
557
|
-
"queue_size": m.queue.qsize(),
|
|
558
|
-
"task": m.current_task or "N/A"
|
|
559
|
-
})
|
|
560
|
-
return out
|
|
561
|
-
except Exception as e:
|
|
562
|
-
ASCIIColors.error(f"Failed to list running models: {e}")
|
|
563
|
-
return []
|
|
564
|
-
|
|
565
|
-
def _acquire_manager(self):
|
|
566
|
-
if self.manager:
|
|
567
|
-
self.registry.release_manager(self.manager.config)
|
|
568
|
-
self.manager = self.registry.get_manager(self.config, self.models_path)
|
|
569
|
-
ASCIIColors.info(f"Binding instance acquired manager for '{self.config['model_name']}'.")
|
|
570
|
-
|
|
571
|
-
def _resolve_device_and_dtype(self):
|
|
572
|
-
if self.config["device"].lower() == "auto":
|
|
573
|
-
self.config["device"] = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
574
|
-
if self.config["torch_dtype_str"].lower() == "auto":
|
|
575
|
-
self.config["torch_dtype_str"] = "float16" if self.config["device"] != "cpu" else "float32"
|
|
576
|
-
|
|
577
|
-
def _decode_image_input(self, item: str) -> Image.Image:
|
|
578
|
-
s = item.strip()
|
|
579
|
-
if s.startswith("data:image/") and ";base64," in s:
|
|
580
|
-
b64 = s.split(";base64,")[-1]
|
|
581
|
-
raw = base64.b64decode(b64)
|
|
582
|
-
return Image.open(BytesIO(raw)).convert("RGB")
|
|
583
|
-
if re_b64 := (s[:30].replace("\n","")):
|
|
584
|
-
try:
|
|
585
|
-
raw = base64.b64decode(s, validate=True)
|
|
586
|
-
return Image.open(BytesIO(raw)).convert("RGB")
|
|
587
|
-
except Exception:
|
|
588
|
-
pass
|
|
589
|
-
try:
|
|
590
|
-
return load_image(s).convert("RGB")
|
|
591
|
-
except Exception:
|
|
592
|
-
return Image.open(s).convert("RGB")
|
|
593
|
-
|
|
594
|
-
def _prepare_seed(self, kwargs: Dict[str, Any]) -> Optional[torch.Generator]:
|
|
595
|
-
seed = kwargs.pop("seed", self.config["seed"])
|
|
596
|
-
if seed == -1:
|
|
597
|
-
return None
|
|
598
|
-
return torch.Generator(device=self.config["device"]).manual_seed(seed)
|
|
599
|
-
|
|
600
|
-
def list_safetensor_models(self) -> List[str]:
|
|
601
|
-
if not self.models_path.exists():
|
|
602
|
-
return []
|
|
603
|
-
return sorted([f.name for f in self.models_path.iterdir() if f.is_file() and f.suffix == ".safetensors"])
|
|
604
|
-
|
|
605
|
-
def listModels(self) -> list:
|
|
606
|
-
civitai_list = [
|
|
607
|
-
{'model_name': key, 'display_name': info['display_name'], 'description': info['description'], 'owned_by': info['owned_by']}
|
|
608
|
-
for key, info in CIVITAI_MODELS.items()
|
|
609
|
-
]
|
|
610
|
-
hf_list = [
|
|
611
|
-
{'model_name': m["model_name"], 'display_name': m["display_name"], 'description': m["desc"], 'owned_by': 'HuggingFace', 'family': m["family"]}
|
|
612
|
-
for m in self.HF_DEFAULT_MODELS
|
|
613
|
-
]
|
|
614
|
-
custom_local = []
|
|
615
|
-
civitai_filenames = {info['filename'] for info in CIVITAI_MODELS.values()}
|
|
616
|
-
for filename in self.list_safetensor_models():
|
|
617
|
-
if filename not in civitai_filenames:
|
|
618
|
-
custom_local.append({'model_name': filename, 'display_name': filename, 'description': 'Local safetensors file.', 'owned_by': 'local_user'})
|
|
619
|
-
return hf_list + civitai_list + custom_local
|
|
620
|
-
|
|
621
|
-
def load_model(self):
|
|
622
|
-
ASCIIColors.info("load_model() called. Loading is automatic on first use.")
|
|
623
|
-
if self.model_name and not self.manager:
|
|
624
|
-
self._acquire_manager()
|
|
625
114
|
|
|
626
|
-
|
|
627
|
-
|
|
628
|
-
|
|
629
|
-
|
|
630
|
-
|
|
631
|
-
|
|
632
|
-
|
|
633
|
-
|
|
634
|
-
|
|
635
|
-
|
|
636
|
-
|
|
637
|
-
|
|
638
|
-
pipeline_args = {
|
|
639
|
-
"prompt": prompt,
|
|
640
|
-
"negative_prompt": negative_prompt or self.config.get("negative_prompt", ""),
|
|
641
|
-
"width": width if width is not None else self.config.get("width", 512),
|
|
642
|
-
"height": height if height is not None else self.config.get("height", 512),
|
|
643
|
-
"num_inference_steps": kwargs.pop("num_inference_steps", self.config.get("num_inference_steps",25)),
|
|
644
|
-
"guidance_scale": kwargs.pop("guidance_scale", self.config.get("guidance_scale",6.5)),
|
|
645
|
-
"generator": generator
|
|
646
|
-
}
|
|
647
|
-
pipeline_args.update(kwargs)
|
|
648
|
-
future = Future()
|
|
649
|
-
self.manager.queue.put((future, "text2image", pipeline_args))
|
|
650
|
-
ASCIIColors.info(f"Job (t2i) '{prompt[:50]}...' queued.")
|
|
651
|
-
try:
|
|
652
|
-
return future.result()
|
|
653
|
-
except Exception as e:
|
|
654
|
-
raise Exception(f"Image generation failed: {e}") from e
|
|
655
|
-
|
|
656
|
-
def _encode_image_to_latents(self, pil: Image.Image, width: int, height: int) -> Tuple[torch.Tensor, Tuple[int,int]]:
|
|
657
|
-
pil = pil.convert("RGB").resize((width, height))
|
|
658
|
-
with self.manager.lock:
|
|
659
|
-
self.manager._load_pipeline_for_task("text2image")
|
|
660
|
-
vae = self.manager.pipeline.vae
|
|
661
|
-
img = torch.from_numpy(torch.ByteTensor(bytearray(pil.tobytes())).numpy()).float() # not efficient but avoids np dep
|
|
662
|
-
img = img.view(pil.height, pil.width, 3).permute(2,0,1).unsqueeze(0) / 255.0
|
|
663
|
-
img = (img * 2.0) - 1.0
|
|
664
|
-
img = img.to(self.config["device"], dtype=getattr(torch, self.config["torch_dtype_str"]))
|
|
665
|
-
with torch.no_grad():
|
|
666
|
-
posterior = vae.encode(img)
|
|
667
|
-
latents = posterior.latent_dist.sample()
|
|
668
|
-
sf = getattr(vae.config, "scaling_factor", 0.18215)
|
|
669
|
-
latents = latents * sf
|
|
670
|
-
return latents, (pil.width, pil.height)
|
|
671
|
-
|
|
672
|
-
def edit_image(self,
|
|
673
|
-
images: Union[str, List[str]],
|
|
674
|
-
prompt: str,
|
|
675
|
-
negative_prompt: Optional[str] = "",
|
|
676
|
-
mask: Optional[str] = None,
|
|
677
|
-
width: Optional[int] = None,
|
|
678
|
-
height: Optional[int] = None,
|
|
679
|
-
**kwargs) -> bytes:
|
|
680
|
-
if not self.model_name:
|
|
681
|
-
raise RuntimeError("No model_name configured. Please select a model in settings.")
|
|
682
|
-
if not self.manager:
|
|
683
|
-
self._acquire_manager()
|
|
684
|
-
imgs = [images] if isinstance(images, str) else list(images)
|
|
685
|
-
pil_images = [self._decode_image_input(s) for s in imgs]
|
|
686
|
-
out_w = width if width is not None else self.config["width"]
|
|
687
|
-
out_h = height if height is not None else self.config["height"]
|
|
688
|
-
generator = self._prepare_seed(kwargs)
|
|
689
|
-
steps = kwargs.pop("num_inference_steps", self.config["num_inference_steps"])
|
|
690
|
-
guidance = kwargs.pop("guidance_scale", self.config["guidance_scale"])
|
|
691
|
-
|
|
692
|
-
# Handle multi-image fusion for Qwen-Image-Edit-2509
|
|
693
|
-
if "Qwen-Image-Edit-2509" in self.model_name and len(pil_images) > 1:
|
|
694
|
-
pipeline_args = {
|
|
695
|
-
"image": pil_images,
|
|
696
|
-
"prompt": prompt,
|
|
697
|
-
"negative_prompt": negative_prompt or " ",
|
|
698
|
-
"width": out_w, "height": out_h,
|
|
699
|
-
"num_inference_steps": steps,
|
|
700
|
-
"true_cfg_scale": guidance,
|
|
701
|
-
"generator": generator
|
|
115
|
+
# Standard dependencies
|
|
116
|
+
pm_v.ensure_packages([
|
|
117
|
+
"pillow", "transformers", "safetensors", "requests", "tqdm", "numpy",
|
|
118
|
+
"accelerate", "uvicorn", "fastapi", "python-multipart", "filelock", "ascii_colors"
|
|
119
|
+
])
|
|
120
|
+
|
|
121
|
+
# Git-based diffusers to get the latest version
|
|
122
|
+
pm_v.ensure_packages([
|
|
123
|
+
{
|
|
124
|
+
"name": "diffusers",
|
|
125
|
+
"vcs": "git+https://github.com/huggingface/diffusers.git",
|
|
126
|
+
"condition": ">=0.35.1"
|
|
702
127
|
}
|
|
703
|
-
|
|
704
|
-
|
|
705
|
-
|
|
706
|
-
|
|
707
|
-
return future.result()
|
|
708
|
-
|
|
709
|
-
# Handle inpainting (single image with mask)
|
|
710
|
-
if mask is not None and len(pil_images) == 1:
|
|
128
|
+
])
|
|
129
|
+
|
|
130
|
+
# XFormers (optional but recommended for NVIDIA)
|
|
131
|
+
if torch_index_url: # Only try to install xformers if CUDA is likely present
|
|
711
132
|
try:
|
|
712
|
-
|
|
133
|
+
ASCIIColors.info("Attempting to install xformers for performance optimization...")
|
|
134
|
+
pm_v.ensure_packages(["xformers"], upgrade=True)
|
|
713
135
|
except Exception as e:
|
|
714
|
-
|
|
715
|
-
|
|
716
|
-
|
|
717
|
-
|
|
718
|
-
|
|
719
|
-
|
|
720
|
-
|
|
721
|
-
|
|
722
|
-
|
|
723
|
-
|
|
724
|
-
|
|
725
|
-
|
|
726
|
-
|
|
727
|
-
|
|
728
|
-
|
|
136
|
+
ASCIIColors.warning(f"Could not install xformers. It's optional but recommended for performance on NVIDIA GPUs. Error: {e}")
|
|
137
|
+
|
|
138
|
+
ASCIIColors.green("Server dependencies are satisfied.")
|
|
139
|
+
|
|
140
|
+
def start_server(self):
|
|
141
|
+
"""
|
|
142
|
+
Installs dependencies and launches the FastAPI server as a background subprocess.
|
|
143
|
+
This method should only be called from within a file lock.
|
|
144
|
+
"""
|
|
145
|
+
server_script = self.server_dir / "main.py"
|
|
146
|
+
if not server_script.exists():
|
|
147
|
+
# Fallback for old structure
|
|
148
|
+
server_script = self.binding_root / "server.py"
|
|
149
|
+
if not server_script.exists():
|
|
150
|
+
raise FileNotFoundError(f"Server script not found at {server_script}. Make sure it's in a 'server' subdirectory.")
|
|
151
|
+
if not self.venv_dir.exists():
|
|
152
|
+
self.install_server_dependencies()
|
|
153
|
+
|
|
154
|
+
if sys.platform == "win32":
|
|
155
|
+
python_executable = self.venv_dir / "Scripts" / "python.exe"
|
|
156
|
+
else:
|
|
157
|
+
python_executable = self.venv_dir / "bin" / "python"
|
|
158
|
+
|
|
159
|
+
command = [
|
|
160
|
+
str(python_executable),
|
|
161
|
+
str(server_script),
|
|
162
|
+
"--host", self.host,
|
|
163
|
+
"--port", str(self.port),
|
|
164
|
+
"--models-path", str(self.models_path.resolve()) # Pass models_path to server
|
|
165
|
+
]
|
|
729
166
|
|
|
730
|
-
|
|
731
|
-
|
|
732
|
-
|
|
733
|
-
|
|
167
|
+
# Use DETACHED_PROCESS on Windows to allow the server to run independently of the parent process.
|
|
168
|
+
# On Linux/macOS, the process will be daemonized enough to not be killed with the worker.
|
|
169
|
+
creationflags = subprocess.DETACHED_PROCESS if sys.platform == "win32" else 0
|
|
170
|
+
|
|
171
|
+
self.server_process = subprocess.Popen(command, creationflags=creationflags)
|
|
172
|
+
ASCIIColors.info("Diffusers server process launched in the background.")
|
|
173
|
+
|
|
174
|
+
def _wait_for_server(self, timeout=300):
|
|
175
|
+
"""Waits for the server to become responsive."""
|
|
176
|
+
ASCIIColors.info("Waiting for Diffusers server to become available...")
|
|
177
|
+
start_time = time.time()
|
|
178
|
+
while time.time() - start_time < timeout:
|
|
179
|
+
if self.is_server_running():
|
|
180
|
+
ASCIIColors.green("Diffusers Server is up and running.")
|
|
181
|
+
# Set initial settings from the binding's config, but only if a model is specified.
|
|
182
|
+
if self.config.get("model_name"):
|
|
183
|
+
try:
|
|
184
|
+
ASCIIColors.info(f"Syncing initial client settings to server (model: {self.config['model_name']})...")
|
|
185
|
+
self.set_settings(self.config)
|
|
186
|
+
except Exception as e:
|
|
187
|
+
ASCIIColors.warning(f"Could not sync initial settings to server: {e}")
|
|
188
|
+
else:
|
|
189
|
+
ASCIIColors.warning("Client has no model_name configured, skipping initial settings sync.")
|
|
190
|
+
return
|
|
191
|
+
time.sleep(2)
|
|
192
|
+
raise RuntimeError("Failed to connect to the Diffusers server within the specified timeout.")
|
|
193
|
+
|
|
194
|
+
def _post_request(self, endpoint: str, data: Optional[dict] = None, files: Optional[dict] = None) -> requests.Response:
|
|
195
|
+
"""Helper to make POST requests to the server."""
|
|
196
|
+
try:
|
|
197
|
+
url = f"{self.base_url}{endpoint}"
|
|
198
|
+
response = requests.post(url, json=data, files=files, timeout=3600) # Long timeout for generation
|
|
199
|
+
response.raise_for_status()
|
|
200
|
+
return response
|
|
201
|
+
except requests.exceptions.RequestException as e:
|
|
202
|
+
ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
|
|
203
|
+
ASCIIColors.error(f"Error details: {e}")
|
|
204
|
+
if hasattr(e, 'response') and e.response:
|
|
205
|
+
try:
|
|
206
|
+
ASCIIColors.error(f"Server response: {e.response.json().get('detail', e.response.text)}")
|
|
207
|
+
except json.JSONDecodeError:
|
|
208
|
+
ASCIIColors.error(f"Server raw response: {e.response.text}")
|
|
209
|
+
raise RuntimeError("Communication with the Diffusers server failed.") from e
|
|
734
210
|
|
|
735
|
-
|
|
211
|
+
def _get_request(self, endpoint: str, params: Optional[dict] = None) -> requests.Response:
|
|
212
|
+
"""Helper to make GET requests to the server."""
|
|
736
213
|
try:
|
|
737
|
-
|
|
738
|
-
|
|
739
|
-
|
|
740
|
-
|
|
741
|
-
|
|
742
|
-
|
|
743
|
-
|
|
744
|
-
"guidance_scale": guidance,
|
|
745
|
-
"generator": generator
|
|
746
|
-
}
|
|
747
|
-
pipeline_args.update(kwargs)
|
|
748
|
-
if "Qwen-Image-Edit" in self.model_name:
|
|
749
|
-
pipeline_args["true_cfg_scale"] = pipeline_args.pop("guidance_scale", 7.0)
|
|
750
|
-
if not pipeline_args.get("negative_prompt"): pipeline_args["negative_prompt"] = " "
|
|
751
|
-
|
|
752
|
-
future = Future()
|
|
753
|
-
self.manager.queue.put((future, "image2image", pipeline_args))
|
|
754
|
-
ASCIIColors.info("Job (i2i) queued.")
|
|
755
|
-
return future.result()
|
|
756
|
-
except Exception:
|
|
757
|
-
pass
|
|
214
|
+
url = f"{self.base_url}{endpoint}"
|
|
215
|
+
response = requests.get(url, params=params, timeout=60)
|
|
216
|
+
response.raise_for_status()
|
|
217
|
+
return response
|
|
218
|
+
except requests.exceptions.RequestException as e:
|
|
219
|
+
ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
|
|
220
|
+
raise RuntimeError("Communication with the Diffusers server failed.") from e
|
|
758
221
|
|
|
759
|
-
|
|
222
|
+
def unload_model(self):
|
|
223
|
+
ASCIIColors.info("Requesting server to unload the current model...")
|
|
760
224
|
try:
|
|
761
|
-
|
|
762
|
-
latents, _ = self._encode_image_to_latents(base, out_w, out_h)
|
|
763
|
-
pipeline_args = {
|
|
764
|
-
"prompt": prompt,
|
|
765
|
-
"negative_prompt": negative_prompt or None,
|
|
766
|
-
"latents": latents,
|
|
767
|
-
"num_inference_steps": steps,
|
|
768
|
-
"guidance_scale": guidance,
|
|
769
|
-
"generator": generator,
|
|
770
|
-
"width": out_w, "height": out_h
|
|
771
|
-
}
|
|
772
|
-
pipeline_args.update(kwargs)
|
|
773
|
-
future = Future()
|
|
774
|
-
self.manager.queue.put((future, "text2image", pipeline_args))
|
|
775
|
-
ASCIIColors.info("Job (t2i with init latents) queued.")
|
|
776
|
-
return future.result()
|
|
225
|
+
self._post_request("/unload_model")
|
|
777
226
|
except Exception as e:
|
|
778
|
-
|
|
227
|
+
ASCIIColors.warning(f"Could not send unload request to server: {e}")
|
|
228
|
+
pass
|
|
229
|
+
|
|
230
|
+
def generate_image(self, prompt: str, negative_prompt: str = "", **kwargs) -> bytes:
|
|
231
|
+
response = self._post_request("/generate_image", data={
|
|
232
|
+
"prompt": prompt,
|
|
233
|
+
"negative_prompt": negative_prompt,
|
|
234
|
+
"params": kwargs
|
|
235
|
+
})
|
|
236
|
+
return response.content
|
|
237
|
+
|
|
238
|
+
def edit_image(self, images: Union[str, List[str], "Image.Image", List["Image.Image"]], prompt: str, **kwargs) -> bytes:
|
|
239
|
+
files = {}
|
|
240
|
+
image_paths = []
|
|
241
|
+
|
|
242
|
+
if not isinstance(images, list):
|
|
243
|
+
images = [images]
|
|
244
|
+
|
|
245
|
+
for i, img in enumerate(images):
|
|
246
|
+
if hasattr(img, 'save'): # PIL Image
|
|
247
|
+
buffer = BytesIO()
|
|
248
|
+
img.save(buffer, format="PNG")
|
|
249
|
+
buffer.seek(0)
|
|
250
|
+
files[f"image_{i}"] = (f"image_{i}.png", buffer, "image/png")
|
|
251
|
+
elif isinstance(img, str) and Path(img).is_file():
|
|
252
|
+
# The server will load this path directly
|
|
253
|
+
image_paths.append(img)
|
|
254
|
+
elif isinstance(img, str): # Handle base64 strings
|
|
255
|
+
try:
|
|
256
|
+
# Simple base64 check
|
|
257
|
+
if img.startswith("data:image/") and ";base64," in img:
|
|
258
|
+
b64_data = img.split(";base64,")[1]
|
|
259
|
+
img_bytes = base64.b64decode(b64_data)
|
|
260
|
+
files[f"image_{i}"] = (f"image_{i}.png", img_bytes, "image/png")
|
|
261
|
+
except Exception:
|
|
262
|
+
raise ValueError(f"Unsupported string image format in edit_image: {img[:100]}")
|
|
263
|
+
else:
|
|
264
|
+
raise ValueError(f"Unsupported image type in edit_image: {type(img)}")
|
|
779
265
|
|
|
266
|
+
data_payload = {
|
|
267
|
+
"prompt": prompt,
|
|
268
|
+
"image_paths": image_paths,
|
|
269
|
+
"params": kwargs
|
|
270
|
+
}
|
|
780
271
|
|
|
781
|
-
|
|
782
|
-
|
|
783
|
-
|
|
784
|
-
|
|
785
|
-
|
|
786
|
-
|
|
787
|
-
]
|
|
788
|
-
safetensors = self.list_safetensor_models()
|
|
789
|
-
return sorted(folders + safetensors)
|
|
272
|
+
# FastAPI needs separate form fields for json and files
|
|
273
|
+
response = self._post_request("/edit_image", data={"json_payload": json.dumps(data_payload)}, files=files)
|
|
274
|
+
return response.content
|
|
275
|
+
|
|
276
|
+
def list_models(self) -> List[Dict[str, Any]]:
|
|
277
|
+
return self._get_request("/list_models").json()
|
|
790
278
|
|
|
279
|
+
def list_local_models(self) -> List[str]:
|
|
280
|
+
return self._get_request("/list_local_models").json()
|
|
281
|
+
|
|
791
282
|
def list_available_models(self) -> List[str]:
|
|
792
|
-
|
|
793
|
-
|
|
794
|
-
return sorted(list(set(local_models + discoverable)))
|
|
795
|
-
|
|
283
|
+
return self._get_request("/list_available_models").json()
|
|
284
|
+
|
|
796
285
|
def list_services(self, **kwargs) -> List[Dict[str, str]]:
|
|
797
|
-
|
|
798
|
-
|
|
799
|
-
if not models:
|
|
800
|
-
return [{"name": "diffusers_no_models", "caption": "No models found", "help": f"Place models in '{self.models_path.resolve()}'."}]
|
|
801
|
-
services = []
|
|
802
|
-
for m in models:
|
|
803
|
-
help_text = "Hugging Face model ID"
|
|
804
|
-
if m in local_models:
|
|
805
|
-
help_text = f"Local model from: {self.models_path.resolve()}"
|
|
806
|
-
elif m in CIVITAI_MODELS:
|
|
807
|
-
help_text = f"Civitai model (downloads as {CIVITAI_MODELS[m]['filename']})"
|
|
808
|
-
services.append({"name": m, "caption": f"Diffusers: {m}", "help": help_text})
|
|
809
|
-
return services
|
|
810
|
-
|
|
286
|
+
return self._get_request("/list_models").json()
|
|
287
|
+
|
|
811
288
|
def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
|
|
812
|
-
|
|
813
|
-
return
|
|
814
|
-
{"name": "model_name", "type": "str", "value": self.model_name, "description": "Local, Civitai, or Hugging Face model.", "options": available_models},
|
|
815
|
-
{"name": "unload_inactive_model_after", "type": "int", "value": self.config["unload_inactive_model_after"], "description": "Unload model after X seconds of inactivity (0 to disable)."},
|
|
816
|
-
{"name": "device", "type": "str", "value": self.config["device"], "description": f"Inference device. Resolved: {self.config['device']}", "options": ["auto","cuda","mps","cpu"]},
|
|
817
|
-
{"name": "torch_dtype_str", "type": "str", "value": self.config["torch_dtype_str"], "description": f"Torch dtype. Resolved: {self.config['torch_dtype_str']}", "options": ["auto","float16","bfloat16","float32"]},
|
|
818
|
-
{"name": "hf_variant", "type": "str", "value": self.config["hf_variant"], "description": "HF model variant (e.g., 'fp16')."},
|
|
819
|
-
{"name": "use_safetensors", "type": "bool", "value": self.config["use_safetensors"], "description": "Prefer .safetensors when loading from Hugging Face."},
|
|
820
|
-
{"name": "scheduler_name", "type": "str", "value": self.config["scheduler_name"], "description": "Scheduler for diffusion.", "options": list(SCHEDULER_MAPPING.keys())},
|
|
821
|
-
{"name": "safety_checker_on", "type": "bool", "value": self.config["safety_checker_on"], "description": "Enable the safety checker."},
|
|
822
|
-
{"name": "enable_cpu_offload", "type": "bool", "value": self.config["enable_cpu_offload"], "description": "Enable model CPU offload (saves VRAM, slower)."},
|
|
823
|
-
{"name": "enable_sequential_cpu_offload", "type": "bool", "value": self.config["enable_sequential_cpu_offload"], "description": "Enable sequential CPU offload."},
|
|
824
|
-
{"name": "enable_xformers", "type": "bool", "value": self.config["enable_xformers"], "description": "Enable xFormers memory efficient attention."},
|
|
825
|
-
{"name": "width", "type": "int", "value": self.config["width"], "description": "Default image width."},
|
|
826
|
-
{"name": "height", "type": "int", "value": self.config["height"], "description": "Default image height."},
|
|
827
|
-
{"name": "num_inference_steps", "type": "int", "value": self.config["num_inference_steps"], "description": "Default inference steps."},
|
|
828
|
-
{"name": "guidance_scale", "type": "float", "value": self.config["guidance_scale"], "description": "Default guidance scale (CFG)."},
|
|
829
|
-
{"name": "seed", "type": "int", "value": self.config["seed"], "description": "Default seed (-1 for random)."},
|
|
830
|
-
{"name": "hf_token", "type": "str", "value": self.config["hf_token"], "description": "HF API token (for private/gated models).", "is_secret": True},
|
|
831
|
-
{"name": "hf_cache_path", "type": "str", "value": self.config["hf_cache_path"], "description": "Path to HF cache."},
|
|
832
|
-
{"name": "local_files_only", "type": "bool", "value": self.config["local_files_only"], "description": "Do not download from Hugging Face."}
|
|
833
|
-
]
|
|
289
|
+
# The server holds the state, so we fetch it.
|
|
290
|
+
return self._get_request("/get_settings").json()
|
|
834
291
|
|
|
835
292
|
def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
|
|
836
|
-
|
|
837
|
-
|
|
838
|
-
|
|
839
|
-
|
|
840
|
-
|
|
841
|
-
|
|
842
|
-
|
|
843
|
-
|
|
844
|
-
|
|
845
|
-
|
|
846
|
-
needs_swap = True
|
|
847
|
-
if needs_swap and self.model_name:
|
|
848
|
-
ASCIIColors.info("Critical settings changed. Swapping model manager...")
|
|
849
|
-
self._resolve_device_and_dtype()
|
|
850
|
-
self._acquire_manager()
|
|
851
|
-
if not needs_swap and self.manager:
|
|
852
|
-
self.manager.config.update(parsed)
|
|
853
|
-
if 'scheduler_name' in parsed and self.manager.pipeline:
|
|
854
|
-
with self.manager.lock:
|
|
855
|
-
self.manager._set_scheduler()
|
|
856
|
-
return True
|
|
293
|
+
# Normalize settings from list of dicts to a single dict if needed
|
|
294
|
+
parsed_settings = settings if isinstance(settings, dict) else {s["name"]: s["value"] for s in settings if "name" in s and "value" in s}
|
|
295
|
+
response = self._post_request("/set_settings", data=parsed_settings)
|
|
296
|
+
return response.json().get("success", False)
|
|
297
|
+
|
|
298
|
+
def ps(self) -> List[dict]:
|
|
299
|
+
try:
|
|
300
|
+
return self._get_request("/ps").json()
|
|
301
|
+
except Exception:
|
|
302
|
+
return [{"error": "Could not connect to server to get process status."}]
|
|
857
303
|
|
|
858
304
|
def __del__(self):
|
|
859
|
-
|
|
860
|
-
|
|
861
|
-
|
|
862
|
-
ASCIIColors.magenta("--- Diffusers TTI Binding Test ---")
|
|
863
|
-
if not DIFFUSERS_AVAILABLE:
|
|
864
|
-
ASCIIColors.error("Diffusers not available. Cannot run test.")
|
|
865
|
-
exit(1)
|
|
866
|
-
temp_paths_dir = Path(__file__).parent / "tmp"
|
|
867
|
-
temp_models_path = temp_paths_dir / "models"
|
|
868
|
-
if temp_paths_dir.exists():
|
|
869
|
-
shutil.rmtree(temp_paths_dir)
|
|
870
|
-
temp_models_path.mkdir(parents=True, exist_ok=True)
|
|
871
|
-
try:
|
|
872
|
-
ASCIIColors.cyan("\n--- Test: Loading a small HF model ---")
|
|
873
|
-
cfg = {"models_path": str(temp_models_path), "model_name": "hf-internal-testing/tiny-stable-diffusion-torch"}
|
|
874
|
-
binding = DiffusersTTIBinding_Impl(**cfg)
|
|
875
|
-
img_bytes = binding.generate_image("a tiny robot", width=64, height=64, num_inference_steps=2)
|
|
876
|
-
assert len(img_bytes) > 1000
|
|
877
|
-
ASCIIColors.green("HF t2i generation OK.")
|
|
878
|
-
del binding
|
|
879
|
-
time.sleep(0.1)
|
|
880
|
-
except Exception as e:
|
|
881
|
-
trace_exception(e)
|
|
882
|
-
ASCIIColors.error(f"Diffusers binding test failed: {e}")
|
|
883
|
-
finally:
|
|
884
|
-
ASCIIColors.cyan("\nCleaning up temporary directories...")
|
|
885
|
-
if temp_paths_dir.exists():
|
|
886
|
-
shutil.rmtree(temp_paths_dir)
|
|
887
|
-
ASCIIColors.magenta("--- Diffusers TTI Binding Test Finished ---")
|
|
305
|
+
# The client destructor does not stop the server,
|
|
306
|
+
# as it is a shared resource for all worker processes.
|
|
307
|
+
pass
|