lollms-client 1.6.2__py3-none-any.whl → 1.6.5__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.

Files changed (41) hide show
  1. lollms_client/__init__.py +1 -1
  2. lollms_client/llm_bindings/azure_openai/__init__.py +2 -2
  3. lollms_client/llm_bindings/claude/__init__.py +2 -2
  4. lollms_client/llm_bindings/gemini/__init__.py +2 -2
  5. lollms_client/llm_bindings/grok/__init__.py +2 -2
  6. lollms_client/llm_bindings/groq/__init__.py +2 -2
  7. lollms_client/llm_bindings/hugging_face_inference_api/__init__.py +2 -2
  8. lollms_client/llm_bindings/litellm/__init__.py +1 -1
  9. lollms_client/llm_bindings/llamacpp/__init__.py +2 -2
  10. lollms_client/llm_bindings/lollms/__init__.py +1 -1
  11. lollms_client/llm_bindings/lollms_webui/__init__.py +1 -1
  12. lollms_client/llm_bindings/mistral/__init__.py +2 -2
  13. lollms_client/llm_bindings/novita_ai/__init__.py +2 -2
  14. lollms_client/llm_bindings/ollama/__init__.py +7 -4
  15. lollms_client/llm_bindings/open_router/__init__.py +2 -2
  16. lollms_client/llm_bindings/openai/__init__.py +1 -1
  17. lollms_client/llm_bindings/openllm/__init__.py +2 -2
  18. lollms_client/llm_bindings/openwebui/__init__.py +1 -1
  19. lollms_client/llm_bindings/perplexity/__init__.py +2 -2
  20. lollms_client/llm_bindings/pythonllamacpp/__init__.py +3 -3
  21. lollms_client/llm_bindings/tensor_rt/__init__.py +1 -1
  22. lollms_client/llm_bindings/transformers/__init__.py +4 -4
  23. lollms_client/llm_bindings/vllm/__init__.py +1 -1
  24. lollms_client/lollms_core.py +7 -1443
  25. lollms_client/lollms_llm_binding.py +1 -1
  26. lollms_client/lollms_tti_binding.py +1 -1
  27. lollms_client/tti_bindings/diffusers/__init__.py +320 -853
  28. lollms_client/tti_bindings/diffusers/server/main.py +882 -0
  29. lollms_client/tti_bindings/gemini/__init__.py +179 -239
  30. lollms_client/tti_bindings/leonardo_ai/__init__.py +1 -1
  31. lollms_client/tti_bindings/novita_ai/__init__.py +1 -1
  32. lollms_client/tti_bindings/stability_ai/__init__.py +1 -1
  33. lollms_client/tts_bindings/lollms/__init__.py +6 -1
  34. lollms_client/tts_bindings/piper_tts/__init__.py +1 -1
  35. lollms_client/tts_bindings/xtts/__init__.py +20 -14
  36. lollms_client/tts_bindings/xtts/server/main.py +17 -4
  37. {lollms_client-1.6.2.dist-info → lollms_client-1.6.5.dist-info}/METADATA +2 -2
  38. {lollms_client-1.6.2.dist-info → lollms_client-1.6.5.dist-info}/RECORD +41 -40
  39. {lollms_client-1.6.2.dist-info → lollms_client-1.6.5.dist-info}/WHEEL +0 -0
  40. {lollms_client-1.6.2.dist-info → lollms_client-1.6.5.dist-info}/licenses/LICENSE +0 -0
  41. {lollms_client-1.6.2.dist-info → lollms_client-1.6.5.dist-info}/top_level.txt +0 -0
@@ -1,887 +1,354 @@
1
- # lollms_client/tti_bindings/diffusers/__init__.py
2
1
  import os
3
- import importlib
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
- from tqdm import tqdm
5
+ import subprocess
6
+ import time
16
7
  import json
17
- import shutil
18
- from lollms_client.lollms_tti_binding import LollmsTTIBinding
19
- from ascii_colors import trace_exception, ASCIIColors
20
-
21
- pm.ensure_packages(["torch","torchvision"],index_url="https://download.pytorch.org/whl/cu126")
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 torch
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
- torch = None
46
- AutoPipelineForText2Image = None
47
- AutoPipelineForImage2Image = None
48
- AutoPipelineForInpainting = None
49
- DiffusionPipeline = None
50
- StableDiffusionPipeline = None
51
- QwenImageEditPipeline = None
52
- QwenImageEditPlusPipeline = None
53
- Image = None
54
- load_image = None
55
- DIFFUSERS_AVAILABLE = False
56
-
57
- BindingName = "DiffusersTTIBinding_Impl"
58
-
59
- CIVITAI_MODELS = {
60
- "realistic-vision-v6": {
61
- "display_name": "Realistic Vision V6.0",
62
- "url": "https://civitai.com/api/download/models/501240?type=Model&format=SafeTensor&size=pruned&fp=fp16",
63
- "filename": "realisticVisionV60_v60B1.safetensors",
64
- "description": "Photorealistic SD1.5 checkpoint.",
65
- "owned_by": "civitai"
66
- },
67
- "absolute-reality": {
68
- "display_name": "Absolute Reality",
69
- "url": "https://civitai.com/api/download/models/132760?type=Model&format=SafeTensor&size=pruned&fp=fp16",
70
- "filename": "absolutereality_v181.safetensors",
71
- "description": "General realistic SD1.5.",
72
- "owned_by": "civitai"
73
- },
74
- "dreamshaper-8": {
75
- "display_name": "DreamShaper 8",
76
- "url": "https://civitai.com/api/download/models/128713",
77
- "filename": "dreamshaper_8.safetensors",
78
- "description": "Versatile SD1.5 style model.",
79
- "owned_by": "civitai"
80
- },
81
- "juggernaut-xl": {
82
- "display_name": "Juggernaut XL",
83
- "url": "https://civitai.com/api/download/models/133005",
84
- "filename": "juggernautXL_version6Rundiffusion.safetensors",
85
- "description": "Artistic SDXL.",
86
- "owned_by": "civitai"
87
- },
88
- "lyriel-v1.6": {
89
- "display_name": "Lyriel v1.6",
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_diffusers_venv")
55
+ self.models_path = Path(kwargs.get("models_path", "./data/models/diffusers_models")).resolve()
56
+ self.models_path.mkdir(exist_ok=True, parents=True)
57
+ if self.auto_start_server:
58
+ self.ensure_server_is_running()
59
+
60
+ def is_server_running(self) -> bool:
61
+ """Checks if the server is already running and responsive."""
295
62
  try:
296
- with requests.get(url, stream=True) as r:
297
- r.raise_for_status()
298
- total_size = int(r.headers.get('content-length', 0))
299
- with open(temp_path, 'wb') as f, tqdm(total=total_size, unit='iB', unit_scale=True, desc=f"Downloading {filename}") as bar:
300
- for chunk in r.iter_content(chunk_size=8192):
301
- f.write(chunk)
302
- bar.update(len(chunk))
303
- shutil.move(temp_path, dest_path)
304
- ASCIIColors.green(f"Model '{filename}' downloaded successfully.")
305
- except Exception as e:
306
- if temp_path.exists():
307
- temp_path.unlink()
308
- raise Exception(f"Failed to download model {filename}: {e}") from e
309
-
310
- def _set_scheduler(self):
311
- if not self.pipeline:
312
- return
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())
63
+ response = requests.get(f"{self.base_url}/status", timeout=2)
64
+ if response.status_code == 200 and response.json().get("status") == "running":
65
+ return True
66
+ except requests.exceptions.RequestException:
67
+ return False
68
+ return False
69
+
70
+
71
+ def ensure_server_is_running(self, continue_if_locked: bool = True):
72
+ """
73
+ Ensures the Diffusers server is running. If not, it attempts to start it
74
+ in a process-safe manner using a file lock.
75
+
76
+ Args:
77
+ continue_if_locked (bool): If True, return immediately if another process
78
+ already holds the lock.
79
+ """
80
+ self.server_dir.mkdir(exist_ok=True)
81
+ lock_path = self.models_path / "diffusers_server.lock"
82
+ lock = FileLock(lock_path)
83
+
84
+ ASCIIColors.info("Attempting to start or connect to the Diffusers server...")
338
85
  try:
339
- load_args = {}
340
- if self.config.get("hf_cache_path"):
341
- load_args["cache_dir"] = str(self.config["hf_cache_path"])
342
- if str(model_path).endswith(".safetensors"):
343
- if task == "text2image":
344
- try:
345
- self.pipeline = AutoPipelineForText2Image.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
346
- except AttributeError:
347
- self.pipeline = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
348
- elif task == "image2image":
349
- self.pipeline = AutoPipelineForImage2Image.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
350
- elif task == "inpainting":
351
- self.pipeline = AutoPipelineForInpainting.from_single_file(model_path, torch_dtype=torch_dtype, cache_dir=load_args.get("cache_dir"))
86
+ # Try to acquire lock immediately if continue_if_locked=True
87
+ with lock.acquire(timeout=0 if continue_if_locked else 60):
88
+ if not self.is_server_running():
89
+ ASCIIColors.yellow("Lock acquired. Starting dedicated Diffusers server...")
90
+ self.start_server()
91
+ else:
92
+ ASCIIColors.green("Server was started by another process. Connected successfully.")
93
+ except Timeout:
94
+ if continue_if_locked:
95
+ ASCIIColors.yellow("Lock held by another process. Skipping server startup and continuing execution.")
96
+ return
352
97
  else:
353
- common_args = {
354
- "torch_dtype": torch_dtype,
355
- "use_safetensors": self.config["use_safetensors"],
356
- "token": self.config["hf_token"],
357
- "local_files_only": self.config["local_files_only"]
358
- }
359
- if self.config["hf_variant"]:
360
- common_args["variant"] = self.config["hf_variant"]
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"]:
388
- try:
389
- self.pipeline.enable_xformers_memory_efficient_attention()
390
- except Exception as e:
391
- ASCIIColors.warning(f"Could not enable xFormers: {e}.")
392
- if self.config["enable_cpu_offload"] and self.config["device"] != "cpu":
393
- self.pipeline.enable_model_cpu_offload()
394
- elif self.config["enable_sequential_cpu_offload"] and self.config["device"] != "cpu":
395
- self.pipeline.enable_sequential_cpu_offload()
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():
98
+ ASCIIColors.yellow("Could not acquire lock within timeout. Waiting for server to become available...")
99
+
100
+ self._wait_for_server()
101
+
102
+ def install_server_dependencies(self):
103
+ """
104
+ Installs the server's dependencies into a dedicated virtual environment
105
+ using pipmaster, which handles complex packages like PyTorch.
106
+ """
107
+ ASCIIColors.info(f"Setting up virtual environment in: {self.venv_dir}")
108
+ pm_v = pm.PackageManager(venv_path=str(self.venv_dir))
109
+
110
+ # --- PyTorch Installation ---
111
+ ASCIIColors.info(f"Installing server dependencies")
112
+ pm_v.ensure_packages([
113
+ "requests", "uvicorn", "fastapi", "python-multipart", "filelock"
114
+ ])
115
+ ASCIIColors.info(f"Installing parisneo libraries")
116
+ pm_v.ensure_packages([
117
+ "ascii_colors","pipmaster"
118
+ ])
119
+ ASCIIColors.info(f"Installing misc libraries (numpy, tqdm...)")
120
+ pm_v.ensure_packages([
121
+ "tqdm", "numpy"
122
+ ])
123
+ ASCIIColors.info(f"Installing Pillow")
124
+ pm_v.ensure_packages([
125
+ "pillow"
126
+ ])
127
+
128
+ ASCIIColors.info(f"Installing pytorch")
129
+ torch_index_url = None
130
+ if sys.platform == "win32":
414
131
  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"
132
+ # Use nvidia-smi to detect CUDA
133
+ result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, check=True)
134
+ ASCIIColors.green("NVIDIA GPU detected. Installing CUDA-enabled PyTorch.")
135
+ # Using a common and stable CUDA version. Adjust if needed.
136
+ torch_index_url = "https://download.pytorch.org/whl/cu128"
137
+ except (FileNotFoundError, subprocess.CalledProcessError):
138
+ 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.")
139
+
140
+ # Base packages including torch. pm.ensure_packages handles verbose output.
141
+ pm_v.ensure_packages(["torch", "torchvision"], index_url=torch_index_url)
142
+
143
+ # Standard dependencies
144
+ ASCIIColors.info(f"Installing transformers dependencies")
145
+ pm_v.ensure_packages([
146
+ "transformers", "safetensors", "accelerate"
147
+ ])
148
+ ASCIIColors.info(f"[Optional] Installing xformers")
149
+ try:
150
+ pm_v.ensure_packages([
151
+ "xformers"
152
+ ])
153
+ except:
154
+ pass
155
+ # Git-based diffusers to get the latest version
156
+ ASCIIColors.info(f"Installing diffusers library from github")
157
+ pm_v.ensure_packages([
158
+ {
159
+ "name": "diffusers",
160
+ "vcs": "git+https://github.com/huggingface/diffusers.git",
161
+ "condition": ">=0.35.1"
162
+ }
163
+ ])
164
+
165
+ ASCIIColors.green("Server dependencies are satisfied.")
166
+
167
+ def start_server(self):
168
+ """
169
+ Installs dependencies and launches the FastAPI server as a background subprocess.
170
+ This method should only be called from within a file lock.
171
+ """
172
+ server_script = self.server_dir / "main.py"
173
+ if not server_script.exists():
174
+ # Fallback for old structure
175
+ server_script = self.binding_root / "server.py"
176
+ if not server_script.exists():
177
+ raise FileNotFoundError(f"Server script not found at {server_script}. Make sure it's in a 'server' subdirectory.")
178
+ if not self.venv_dir.exists():
179
+ self.install_server_dependencies()
180
+
181
+ if sys.platform == "win32":
182
+ python_executable = self.venv_dir / "Scripts" / "python.exe"
183
+ else:
184
+ python_executable = self.venv_dir / "bin" / "python"
185
+
186
+ command = [
187
+ str(python_executable),
188
+ str(server_script),
189
+ "--host", self.host,
190
+ "--port", str(self.port),
191
+ "--models-path", str(self.models_path.resolve()) # Pass models_path to server
455
192
  ]
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", "")
523
-
524
- models_path_str = kwargs.get("models_path", str(Path(__file__).parent / "models"))
525
- self.models_path = Path(models_path_str)
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()
531
193
 
532
- def ps(self) -> List[dict]:
533
- if not self.registry:
534
- return []
194
+ # Use DETACHED_PROCESS on Windows to allow the server to run independently of the parent process.
195
+ # On Linux/macOS, the process will be daemonized enough to not be killed with the worker.
196
+ creationflags = subprocess.DETACHED_PROCESS if sys.platform == "win32" else 0
197
+
198
+ self.server_process = subprocess.Popen(command, creationflags=creationflags)
199
+ ASCIIColors.info("Diffusers server process launched in the background.")
200
+
201
+ def _wait_for_server(self, timeout=300):
202
+ """Waits for the server to become responsive."""
203
+ ASCIIColors.info("Waiting for Diffusers server to become available...")
204
+ start_time = time.time()
205
+ while time.time() - start_time < timeout:
206
+ if self.is_server_running():
207
+ ASCIIColors.green("Diffusers Server is up and running.")
208
+ # Set initial settings from the binding's config, but only if a model is specified.
209
+ if self.config.get("model_name"):
210
+ try:
211
+ ASCIIColors.info(f"Syncing initial client settings to server (model: {self.config['model_name']})...")
212
+ self.set_settings(self.config)
213
+ except Exception as e:
214
+ ASCIIColors.warning(f"Could not sync initial settings to server: {e}")
215
+ else:
216
+ ASCIIColors.warning("Client has no model_name configured, skipping initial settings sync.")
217
+ return
218
+ time.sleep(2)
219
+ raise RuntimeError("Failed to connect to the Diffusers server within the specified timeout.")
220
+
221
+ def _post_json_request(self, endpoint: str, data: Optional[dict] = None) -> requests.Response:
222
+ """Helper to make POST requests with a JSON body."""
535
223
  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
224
+ url = f"{self.base_url}{endpoint}"
225
+ response = requests.post(url, json=data, timeout=3600) # Long timeout for generation
226
+ response.raise_for_status()
227
+ return response
228
+ except requests.exceptions.RequestException as e:
229
+ ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
230
+ ASCIIColors.error(f"Error details: {e}")
231
+ if hasattr(e, 'response') and e.response:
232
+ try:
233
+ ASCIIColors.error(f"Server response: {e.response.json().get('detail', e.response.text)}")
234
+ except json.JSONDecodeError:
235
+ ASCIIColors.error(f"Server raw response: {e.response.text}")
236
+ raise RuntimeError("Communication with the Diffusers server failed.") from e
237
+
238
+ def _post_multipart_request(self, endpoint: str, data: Optional[dict] = None, files: Optional[list] = None) -> requests.Response:
239
+ """Helper to make multipart/form-data POST requests for file uploads."""
589
240
  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()
241
+ url = f"{self.base_url}{endpoint}"
242
+ response = requests.post(url, data=data, files=files, timeout=3600)
243
+ response.raise_for_status()
244
+ return response
245
+ except requests.exceptions.RequestException as e:
246
+ # (Error handling is the same as above)
247
+ ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
248
+ ASCIIColors.error(f"Error details: {e}")
249
+ if hasattr(e, 'response') and e.response:
250
+ try:
251
+ ASCIIColors.error(f"Server response: {e.response.json().get('detail', e.response.text)}")
252
+ except json.JSONDecodeError:
253
+ ASCIIColors.error(f"Server raw response: {e.response.text}")
254
+ raise RuntimeError("Communication with the Diffusers server failed.") from e
255
+
256
+ def _get_request(self, endpoint: str, params: Optional[dict] = None) -> requests.Response:
257
+ """Helper to make GET requests to the server."""
258
+ try:
259
+ url = f"{self.base_url}{endpoint}"
260
+ response = requests.get(url, params=params, timeout=60)
261
+ response.raise_for_status()
262
+ return response
263
+ except requests.exceptions.RequestException as e:
264
+ ASCIIColors.error(f"Failed to communicate with Diffusers server at {url}.")
265
+ raise RuntimeError("Communication with the Diffusers server failed.") from e
625
266
 
626
267
  def unload_model(self):
627
- if self.manager:
628
- ASCIIColors.info(f"Binding instance releasing manager for '{self.manager.config['model_name']}'.")
629
- self.registry.release_manager(self.manager.config)
630
- self.manager = None
631
-
632
- def generate_image(self, prompt: str, negative_prompt: str = "", width: int|None = None, height: int|None = None, **kwargs) -> bytes:
633
- if not self.model_name:
634
- raise RuntimeError("No model_name configured. Please select a model in settings.")
635
- if not self.manager:
636
- self._acquire_manager()
637
- generator = self._prepare_seed(kwargs)
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.")
268
+ ASCIIColors.info("Requesting server to unload the current model...")
651
269
  try:
652
- return future.result()
270
+ self._post_json_request("/unload_model")
653
271
  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
702
- }
703
- pipeline_args.update(kwargs)
704
- future = Future()
705
- self.manager.queue.put((future, "image2image", pipeline_args))
706
- ASCIIColors.info(f"Job (multi-image fusion with {len(pil_images)} images) queued.")
707
- return future.result()
708
-
709
- # Handle inpainting (single image with mask)
710
- if mask is not None and len(pil_images) == 1:
711
- try:
712
- mask_img = self._decode_image_input(mask).convert("L")
713
- except Exception as e:
714
- raise ValueError(f"Failed to decode mask image: {e}") from e
715
- pipeline_args = {
716
- "image": pil_images[0],
717
- "mask_image": mask_img,
718
- "prompt": prompt,
719
- "negative_prompt": negative_prompt or None,
720
- "width": out_w, "height": out_h,
721
- "num_inference_steps": steps,
722
- "guidance_scale": guidance,
723
- "generator": generator
724
- }
725
- pipeline_args.update(kwargs)
726
- if "Qwen-Image-Edit" in self.model_name:
727
- pipeline_args["true_cfg_scale"] = pipeline_args.pop("guidance_scale", 7.0)
728
- if not pipeline_args.get("negative_prompt"): pipeline_args["negative_prompt"] = " "
272
+ ASCIIColors.warning(f"Could not send unload request to server: {e}")
273
+ pass
729
274
 
730
- future = Future()
731
- self.manager.queue.put((future, "inpainting", pipeline_args))
732
- ASCIIColors.info("Job (inpaint) queued.")
733
- return future.result()
275
+ def generate_image(self, prompt: str, negative_prompt: str = "", **kwargs) -> bytes:
276
+ # This is a pure JSON request
277
+ response = self._post_json_request("/generate_image", data={
278
+ "prompt": prompt,
279
+ "negative_prompt": negative_prompt,
280
+ "params": kwargs
281
+ })
282
+ return response.content
283
+
284
+ def edit_image(self, images: Union[str, List[str], "Image.Image", List["Image.Image"]], prompt: str, **kwargs) -> bytes:
285
+ images_b64 = []
286
+ if not isinstance(images, list):
287
+ images = [images]
288
+
289
+
290
+ for img in images:
291
+ # Case 1: Input is a PIL Image object
292
+ if hasattr(img, 'save'):
293
+ buffer = BytesIO()
294
+ img.save(buffer, format="PNG")
295
+ b64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
296
+ images_b64.append(b64_string)
297
+
298
+ # Case 2: Input is a string (could be path or already base64)
299
+ elif isinstance(img, str):
300
+ try:
301
+ b64_string = img.split(";base64,")[1] if ";base64," in img else img
302
+ base64.b64decode(b64_string) # Validate
303
+ images_b64.append(b64_string)
304
+ except Exception:
305
+ ASCIIColors.warning(f"Warning: A string input was not a valid file path or base64. Skipping.")
306
+ else:
307
+ raise ValueError(f"Unsupported image type in edit_image: {type(img)}")
308
+ if not images_b64:
309
+ raise ValueError("No valid images were provided to the edit_image function.")
734
310
 
735
- # Handle standard image-to-image (single image)
736
- try:
737
- pipeline_args = {
738
- "image": pil_images[0],
739
- "prompt": prompt,
740
- "negative_prompt": negative_prompt or None,
741
- "strength": kwargs.pop("strength", 0.6),
742
- "width": out_w, "height": out_h,
743
- "num_inference_steps": steps,
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
311
+ # Translate "mask" to "mask_image" for server compatibility
312
+ if "mask" in kwargs and kwargs["mask"]:
313
+ kwargs["mask_image"] = kwargs.pop("mask")
758
314
 
759
- # Fallback to latent-based generation if i2i fails for some reason
760
- try:
761
- base = pil_images[0]
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()
777
- except Exception as e:
778
- raise Exception(f"Image edit failed: {e}") from e
315
+ json_payload = {
316
+ "prompt": prompt,
317
+ "images_b64": images_b64,
318
+ "params": kwargs
319
+ }
320
+ response = self._post_json_request("/edit_image", data=json_payload)
321
+ return response.content
779
322
 
323
+ def list_models(self) -> List[Dict[str, Any]]:
324
+ return self._get_request("/list_models").json()
780
325
 
781
326
  def list_local_models(self) -> List[str]:
782
- if not self.models_path.exists():
783
- return []
784
- folders = [
785
- d.name for d in self.models_path.iterdir()
786
- if d.is_dir() and ((d / "model_index.json").exists() or (d / "unet" / "config.json").exists())
787
- ]
788
- safetensors = self.list_safetensor_models()
789
- return sorted(folders + safetensors)
327
+ return self._get_request("/list_local_models").json()
790
328
 
791
329
  def list_available_models(self) -> List[str]:
792
- discoverable = [m['model_name'] for m in self.listModels()]
793
- local_models = self.list_local_models()
794
- return sorted(list(set(local_models + discoverable)))
330
+ return self._get_request("/list_available_models").json()
795
331
 
796
332
  def list_services(self, **kwargs) -> List[Dict[str, str]]:
797
- models = self.list_available_models()
798
- local_models = self.list_local_models()
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
333
+ return self._get_request("/list_models").json()
810
334
 
811
335
  def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
812
- available_models = self.list_available_models()
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
- ]
336
+ # The server holds the state, so we fetch it.
337
+ return self._get_request("/get_settings").json()
834
338
 
835
339
  def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
836
- parsed = settings if isinstance(settings, dict) else {i["name"]: i["value"] for i in settings if "name" in i and "value" in i}
837
- critical_keys = self.registry._get_critical_keys()
838
- needs_swap = False
839
- for key, value in parsed.items():
840
- if self.config.get(key) != value:
841
- ASCIIColors.info(f"Setting '{key}' changed to: {value}")
842
- self.config[key] = value
843
- if key == "model_name":
844
- self.model_name = value
845
- if key in critical_keys:
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
340
+ # Normalize settings from list of dicts to a single dict if needed
341
+ parsed_settings = settings if isinstance(settings, dict) else {s["name"]: s["value"] for s in settings if "name" in s and "value" in s}
342
+ response = self._post_json_request("/set_settings", data=parsed_settings)
343
+ return response.json().get("success", False)
344
+
345
+ def ps(self) -> List[dict]:
346
+ try:
347
+ return self._get_request("/ps").json()
348
+ except Exception:
349
+ return [{"error": "Could not connect to server to get process status."}]
857
350
 
858
351
  def __del__(self):
859
- self.unload_model()
860
-
861
- if __name__ == '__main__':
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 ---")
352
+ # The client destructor does not stop the server,
353
+ # as it is a shared resource for all worker processes.
354
+ pass