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

@@ -0,0 +1,692 @@
1
+ # lollms_client/tti_bindings/diffusers/__init__.py
2
+ import os
3
+ import importlib
4
+ from io import BytesIO
5
+ from typing import Optional, List, Dict, Any, Union
6
+ from pathlib import Path
7
+
8
+
9
+ try:
10
+ import pipmaster as pm
11
+ import platform # For OS detection for torch index
12
+
13
+ # Determine initial device preference to guide torch installation
14
+ preferred_torch_device_for_install = "cpu" # Default assumption
15
+
16
+ # Tentatively set preference based on OS, assuming user might want GPU if available
17
+ if platform.system() == "Linux" or platform.system() == "Windows":
18
+ # On Linux/Windows, CUDA is the primary GPU acceleration for PyTorch.
19
+ # We will try to install a CUDA version of PyTorch.
20
+ preferred_torch_device_for_install = "cuda"
21
+ elif platform.system() == "Darwin":
22
+ # On macOS, MPS is the acceleration. Standard torch install usually handles this.
23
+ preferred_torch_device_for_install = "mps" # or keep cpu if mps detection is later
24
+
25
+ torch_pkgs = ["torch", "torchaudio", "torchvision", "xformers"]
26
+ diffusers_core_pkgs = ["diffusers", "Pillow", "transformers", "safetensors"]
27
+
28
+ torch_index_url = None
29
+ if preferred_torch_device_for_install == "cuda":
30
+ # Specify a common CUDA version index. Pip should resolve the correct torch version.
31
+ # As of late 2023/early 2024, cu118 or cu121 are common. Let's use cu126.
32
+ # Users with different CUDA setups might need to pre-install torch manually.
33
+ torch_index_url = "https://download.pytorch.org/whl/cu126"
34
+ ASCIIColors.info(f"Attempting to ensure PyTorch with CUDA support (target index: {torch_index_url})")
35
+ # Install torch and torchaudio first from the specific index
36
+ pm.ensure_packages(torch_pkgs, index_url=torch_index_url)
37
+ # Then install audiocraft and other dependencies; pip should use the already installed torch
38
+ pm.ensure_packages(diffusers_core_pkgs)
39
+ else:
40
+ # For CPU, MPS, or if no specific CUDA preference was determined for install
41
+ ASCIIColors.info("Ensuring PyTorch, AudioCraft, and dependencies using default PyPI index.")
42
+ pm.ensure_packages(torch_pkgs + diffusers_core_pkgs)
43
+
44
+ import whisper
45
+ import torch
46
+ _whisper_installed = True
47
+ except Exception as e:
48
+ _whisper_installation_error = str(e)
49
+ whisper = None
50
+ torch = None
51
+
52
+
53
+ try:
54
+ import torch
55
+ from diffusers import AutoPipelineForText2Image, DiffusionPipeline
56
+ from diffusers.utils import load_image # Potentially for future img2img etc.
57
+ from PIL import Image
58
+ DIFFUSERS_AVAILABLE = True
59
+ except ImportError:
60
+ torch = None
61
+ AutoPipelineForText2Image = None
62
+ DiffusionPipeline = None
63
+ Image = None
64
+ load_image = None
65
+ DIFFUSERS_AVAILABLE = False
66
+ # Detailed error will be raised in __init__ if user tries to use it
67
+
68
+ from lollms_client.lollms_tti_binding import LollmsTTIBinding
69
+ from ascii_colors import trace_exception, ASCIIColors
70
+ import json # For potential JSONDecodeError and settings
71
+
72
+ # Defines the binding name for the manager
73
+ BindingName = "DiffusersTTIBinding_Impl"
74
+
75
+ # Helper for torch.dtype string conversion
76
+ TORCH_DTYPE_MAP_STR_TO_OBJ = {
77
+ "float16": torch.float16 if torch else "float16", # Keep string if torch not loaded
78
+ "bfloat16": torch.bfloat16 if torch else "bfloat16",
79
+ "float32": torch.float32 if torch else "float32",
80
+ "auto": "auto"
81
+ }
82
+ TORCH_DTYPE_MAP_OBJ_TO_STR = {v: k for k, v in TORCH_DTYPE_MAP_STR_TO_OBJ.items()}
83
+ if torch: # Add None mapping if torch is loaded
84
+ TORCH_DTYPE_MAP_OBJ_TO_STR[None] = "None"
85
+
86
+
87
+ # Common Schedulers mapping (User-friendly name to Class name)
88
+ SCHEDULER_MAPPING = {
89
+ "default": None, # Use model's default
90
+ "ddim": "DDIMScheduler",
91
+ "ddpm": "DDPMScheduler",
92
+ "deis_multistep": "DEISMultistepScheduler",
93
+ "dpm_multistep": "DPMSolverMultistepScheduler", # Alias
94
+ "dpm_multistep_karras": "DPMSolverMultistepScheduler", # Configured with use_karras_sigmas=True
95
+ "dpm_single": "DPMSolverSinglestepScheduler",
96
+ "dpm_adaptive": "DP soluzioniPlusPlusScheduler", # DPM++ 2M Karras in A1111
97
+ "dpm++_2m": "DPMSolverMultistepScheduler",
98
+ "dpm++_2m_karras": "DPMSolverMultistepScheduler", # Configured with use_karras_sigmas=True
99
+ "dpm++_2s_ancestral": "DPMSolverAncestralDiscreteScheduler",
100
+ "dpm++_2s_ancestral_karras": "DPMSolverAncestralDiscreteScheduler", # Configured with use_karras_sigmas=True
101
+ "dpm++_sde": "DPMSolverSDEScheduler",
102
+ "dpm++_sde_karras": "DPMSolverSDEScheduler", # Configured with use_karras_sigmas=True
103
+ "euler_ancestral_discrete": "EulerAncestralDiscreteScheduler",
104
+ "euler_discrete": "EulerDiscreteScheduler",
105
+ "heun_discrete": "HeunDiscreteScheduler",
106
+ "heun_karras": "HeunDiscreteScheduler", # Configured with use_karras_sigmas=True
107
+ "lms_discrete": "LMSDiscreteScheduler",
108
+ "lms_karras": "LMSDiscreteScheduler", # Configured with use_karras_sigmas=True
109
+ "pndm": "PNDMScheduler",
110
+ "unipc_multistep": "UniPCMultistepScheduler",
111
+ }
112
+ SCHEDULER_USES_KARRAS_SIGMAS = [
113
+ "dpm_multistep_karras", "dpm++_2m_karras", "dpm++_2s_ancestral_karras",
114
+ "dpm++_sde_karras", "heun_karras", "lms_karras"
115
+ ]
116
+
117
+
118
+ class DiffusersTTIBinding_Impl(LollmsTTIBinding):
119
+ """
120
+ Concrete implementation of LollmsTTIBinding for Hugging Face Diffusers library.
121
+ Allows running various text-to-image models locally.
122
+ """
123
+ DEFAULT_CONFIG = {
124
+ "model_id_or_path": "stabilityai/stable-diffusion-2-1-base",
125
+ "device": "auto", # "auto", "cuda", "mps", "cpu"
126
+ "torch_dtype_str": "auto", # "auto", "float16", "bfloat16", "float32"
127
+ "use_safetensors": True,
128
+ "scheduler_name": "default",
129
+ "safety_checker_on": True, # Note: Diffusers default is ON
130
+ "num_inference_steps": 25,
131
+ "guidance_scale": 7.5,
132
+ "default_width": 768, # Default for SD 2.1 base
133
+ "default_height": 768, # Default for SD 2.1 base
134
+ "seed": -1, # -1 for random on each call
135
+ "enable_cpu_offload": False,
136
+ "enable_sequential_cpu_offload": False,
137
+ "enable_xformers": False, # Explicit opt-in for xformers
138
+ "hf_variant": None, # e.g., "fp16"
139
+ "hf_token": None,
140
+ "local_files_only": False,
141
+ }
142
+
143
+
144
+ def __init__(self,
145
+ config: Optional[Dict[str, Any]] = None,
146
+ lollms_paths: Optional[Dict[str, Union[str, Path]]] = None,
147
+ **kwargs # Catches other potential parameters like 'service_key' or 'client_id'
148
+ ):
149
+ """
150
+ Initialize the Diffusers TTI binding.
151
+
152
+ Args:
153
+ config (Optional[Dict[str, Any]]): Configuration dictionary for the binding.
154
+ Overrides DEFAULT_CONFIG.
155
+ lollms_paths (Optional[Dict[str, Union[str, Path]]]): Dictionary of LOLLMS paths.
156
+ Used for model/cache directories.
157
+ **kwargs: Catches other parameters (e.g. service_key).
158
+ """
159
+ super().__init__(binding_name="diffusers")
160
+
161
+ if not DIFFUSERS_AVAILABLE:
162
+ ASCIIColors.error("Diffusers library or its dependencies (torch, Pillow, transformers) are not installed or failed to import.")
163
+ ASCIIColors.info("Attempting to install/verify packages...")
164
+ pm.ensure_packages(["torch", "diffusers", "Pillow", "transformers", "safetensors"])
165
+ try:
166
+ import torch as _torch
167
+ from diffusers import AutoPipelineForText2Image as _AutoPipelineForText2Image
168
+ from diffusers import DiffusionPipeline as _DiffusionPipeline
169
+ from PIL import Image as _Image
170
+ globals()['torch'] = _torch
171
+ globals()['AutoPipelineForText2Image'] = _AutoPipelineForText2Image
172
+ globals()['DiffusionPipeline'] = _DiffusionPipeline
173
+ globals()['Image'] = _Image
174
+
175
+ # Re-populate torch dtype maps if torch was just loaded
176
+ global TORCH_DTYPE_MAP_STR_TO_OBJ, TORCH_DTYPE_MAP_OBJ_TO_STR
177
+ TORCH_DTYPE_MAP_STR_TO_OBJ = {
178
+ "float16": _torch.float16, "bfloat16": _torch.bfloat16,
179
+ "float32": _torch.float32, "auto": "auto"
180
+ }
181
+ TORCH_DTYPE_MAP_OBJ_TO_STR = {v: k for k, v in TORCH_DTYPE_MAP_STR_TO_OBJ.items()}
182
+ TORCH_DTYPE_MAP_OBJ_TO_STR[None] = "None"
183
+ ASCIIColors.green("Dependencies seem to be available now.")
184
+ except ImportError as e:
185
+ trace_exception(e)
186
+ raise ImportError(
187
+ "Diffusers binding dependencies are still not met after trying to ensure them. "
188
+ "Please install torch, diffusers, Pillow, and transformers manually. "
189
+ f"Error: {e}"
190
+ ) from e
191
+
192
+ self.config = {**self.DEFAULT_CONFIG, **(config or {}), **kwargs}
193
+ self.lollms_paths = {k: Path(v) for k, v in lollms_paths.items()} if lollms_paths else {}
194
+
195
+ self.pipeline: Optional[DiffusionPipeline] = None
196
+ self.current_model_id_or_path = None # To track if model needs reload
197
+
198
+ # Resolve auto settings for device and dtype
199
+ if self.config["device"].lower() == "auto":
200
+ if torch.cuda.is_available(): self.config["device"] = "cuda"
201
+ elif torch.backends.mps.is_available(): self.config["device"] = "mps"
202
+ else: self.config["device"] = "cpu"
203
+
204
+ if self.config["torch_dtype_str"].lower() == "auto":
205
+ if self.config["device"] == "cpu": self.config["torch_dtype_str"] = "float32" # CPU usually float32
206
+ else: self.config["torch_dtype_str"] = "float16" # Common default for GPU
207
+
208
+ self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
209
+ if self.torch_dtype == "auto": # Should have been resolved above
210
+ self.torch_dtype = torch.float16 if self.config["device"] != "cpu" else torch.float32
211
+ self.config["torch_dtype_str"] = TORCH_DTYPE_MAP_OBJ_TO_STR.get(self.torch_dtype, "float32")
212
+
213
+
214
+ # For potential lollms client specific features
215
+ self.client_id = kwargs.get("service_key", kwargs.get("client_id", "diffusers_client_user"))
216
+
217
+ self.load_model()
218
+
219
+
220
+ def _resolve_model_path(self, model_id_or_path: str) -> str:
221
+ """Resolves a model name/path against lollms_paths if not absolute."""
222
+ if os.path.isabs(model_id_or_path):
223
+ return model_id_or_path
224
+
225
+ # Check personal_models_path/diffusers_models/<name>
226
+ if self.lollms_paths.get('personal_models_path'):
227
+ personal_diffusers_path = self.lollms_paths['personal_models_path'] / "diffusers_models" / model_id_or_path
228
+ if personal_diffusers_path.exists() and personal_diffusers_path.is_dir():
229
+ ASCIIColors.info(f"Found local model in personal_models_path: {personal_diffusers_path}")
230
+ return str(personal_diffusers_path)
231
+
232
+ # Check models_zoo_path/diffusers_models/<name> (if different from personal)
233
+ if self.lollms_paths.get('models_zoo_path') and \
234
+ self.lollms_paths.get('models_zoo_path') != self.lollms_paths.get('personal_models_path'):
235
+ zoo_diffusers_path = self.lollms_paths['models_zoo_path'] / "diffusers_models" / model_id_or_path
236
+ if zoo_diffusers_path.exists() and zoo_diffusers_path.is_dir():
237
+ ASCIIColors.info(f"Found local model in models_zoo_path: {zoo_diffusers_path}")
238
+ return str(zoo_diffusers_path)
239
+
240
+ ASCIIColors.info(f"Assuming '{model_id_or_path}' is a Hugging Face Hub ID or already fully qualified.")
241
+ return model_id_or_path
242
+
243
+ def load_model(self):
244
+ """Loads the Diffusers pipeline based on current configuration."""
245
+ ASCIIColors.info("Loading Diffusers model...")
246
+ if self.pipeline is not None:
247
+ self.unload_model() # Ensure old model is cleared
248
+
249
+ try:
250
+ model_path = self._resolve_model_path(self.config["model_id_or_path"])
251
+ self.current_model_id_or_path = model_path # Store what's actually loaded
252
+
253
+ load_args = {
254
+ "torch_dtype": self.torch_dtype,
255
+ "use_safetensors": self.config["use_safetensors"],
256
+ "token": self.config["hf_token"],
257
+ "local_files_only": self.config["local_files_only"],
258
+ }
259
+ if self.config["hf_variant"]:
260
+ load_args["variant"] = self.config["hf_variant"]
261
+
262
+ if not self.config["safety_checker_on"]:
263
+ load_args["safety_checker"] = None
264
+
265
+ if self.lollms_paths.get("shared_cache_path"):
266
+ load_args["cache_dir"] = str(self.lollms_paths["shared_cache_path"] / "huggingface_diffusers")
267
+
268
+
269
+ # Use AutoPipelineForText2Image for flexibility
270
+ pipeline_class_to_load = AutoPipelineForText2Image
271
+ custom_pipeline_class_name = self.config.get("pipeline_class_name")
272
+
273
+ if custom_pipeline_class_name:
274
+ try:
275
+ diffusers_module = importlib.import_module("diffusers")
276
+ pipeline_class_to_load = getattr(diffusers_module, custom_pipeline_class_name)
277
+ ASCIIColors.info(f"Using specified pipeline class: {custom_pipeline_class_name}")
278
+ except (ImportError, AttributeError) as e:
279
+ ASCIIColors.warning(f"Could not load custom pipeline class {custom_pipeline_class_name}: {e}. Falling back to AutoPipelineForText2Image.")
280
+ pipeline_class_to_load = AutoPipelineForText2Image
281
+
282
+ self.pipeline = pipeline_class_to_load.from_pretrained(model_path, **load_args)
283
+
284
+ # Scheduler
285
+ self._set_scheduler()
286
+
287
+ self.pipeline.to(self.config["device"])
288
+
289
+ if self.config["enable_xformers"]:
290
+ try:
291
+ self.pipeline.enable_xformers_memory_efficient_attention()
292
+ ASCIIColors.info("xFormers memory efficient attention enabled.")
293
+ except Exception as e:
294
+ ASCIIColors.warning(f"Could not enable xFormers: {e}. Proceeding without it.")
295
+
296
+ if self.config["enable_cpu_offload"] and self.config["device"] != "cpu":
297
+ self.pipeline.enable_model_cpu_offload()
298
+ ASCIIColors.info("Model CPU offload enabled.")
299
+ elif self.config["enable_sequential_cpu_offload"] and self.config["device"] != "cpu":
300
+ self.pipeline.enable_sequential_cpu_offload() # More aggressive
301
+ ASCIIColors.info("Sequential CPU offload enabled.")
302
+
303
+
304
+ ASCIIColors.green(f"Diffusers model '{model_path}' loaded successfully on device '{self.config['device']}' with dtype '{self.config['torch_dtype_str']}'.")
305
+
306
+ except Exception as e:
307
+ trace_exception(e)
308
+ self.pipeline = None
309
+ raise RuntimeError(f"Failed to load Diffusers model '{self.config['model_id_or_path']}': {e}") from e
310
+
311
+ def _set_scheduler(self):
312
+ if not self.pipeline: return
313
+
314
+ scheduler_name_key = self.config["scheduler_name"].lower()
315
+ if scheduler_name_key == "default":
316
+ ASCIIColors.info(f"Using model's default scheduler: {self.pipeline.scheduler.__class__.__name__}")
317
+ return
318
+
319
+ scheduler_class_name = SCHEDULER_MAPPING.get(scheduler_name_key)
320
+ if scheduler_class_name:
321
+ try:
322
+ scheduler_module = importlib.import_module("diffusers.schedulers")
323
+ SchedulerClass = getattr(scheduler_module, scheduler_class_name)
324
+
325
+ scheduler_config = self.pipeline.scheduler.config
326
+ if scheduler_name_key in SCHEDULER_USES_KARRAS_SIGMAS:
327
+ scheduler_config["use_karras_sigmas"] = True
328
+ else: # Ensure it's False if not a karras variant for this scheduler
329
+ if "use_karras_sigmas" in scheduler_config:
330
+ scheduler_config["use_karras_sigmas"] = False
331
+
332
+
333
+ self.pipeline.scheduler = SchedulerClass.from_config(scheduler_config)
334
+ ASCIIColors.info(f"Switched scheduler to {scheduler_name_key} ({scheduler_class_name}).")
335
+ except Exception as e:
336
+ trace_exception(e)
337
+ ASCIIColors.warning(f"Could not switch scheduler to {scheduler_name_key}: {e}. Using current default.")
338
+ else:
339
+ ASCIIColors.warning(f"Unknown scheduler name: {self.config['scheduler_name']}. Using model default.")
340
+
341
+
342
+ def unload_model(self):
343
+ if self.pipeline is not None:
344
+ del self.pipeline
345
+ self.pipeline = None
346
+ ASCIIColors.info("Diffusers pipeline unloaded.")
347
+ if torch and torch.cuda.is_available():
348
+ torch.cuda.empty_cache()
349
+
350
+ def generate_image(self,
351
+ prompt: str,
352
+ negative_prompt: Optional[str] = "",
353
+ width: Optional[int] = None, # Uses default from config if None
354
+ height: Optional[int] = None, # Uses default from config if None
355
+ **kwargs) -> bytes:
356
+ """
357
+ Generates image data using the Diffusers pipeline.
358
+
359
+ Args:
360
+ prompt (str): The positive text prompt.
361
+ negative_prompt (Optional[str]): The negative prompt.
362
+ width (int): Image width. Overrides default.
363
+ height (int): Image height. Overrides default.
364
+ **kwargs: Additional parameters for the pipeline:
365
+ - num_inference_steps (int)
366
+ - guidance_scale (float)
367
+ - seed (int)
368
+ - eta (float, for DDIM)
369
+ - num_images_per_prompt (int, though this binding returns one)
370
+ - clip_skip (int, if supported by pipeline - advanced)
371
+ Returns:
372
+ bytes: The generated image data (PNG format).
373
+ Raises:
374
+ Exception: If the request fails or image generation fails.
375
+ """
376
+ if not self.pipeline:
377
+ raise RuntimeError("Diffusers pipeline is not loaded. Cannot generate image.")
378
+
379
+ # Use call-specific or configured defaults
380
+ _width = width if width is not None else self.config["default_width"]
381
+ _height = height if height is not None else self.config["default_height"]
382
+ _num_inference_steps = kwargs.get("num_inference_steps", self.config["num_inference_steps"])
383
+ _guidance_scale = kwargs.get("guidance_scale", self.config["guidance_scale"])
384
+ _seed = kwargs.get("seed", self.config["seed"])
385
+
386
+ generator = None
387
+ if _seed != -1: # -1 means random seed
388
+ generator = torch.Generator(device=self.config["device"]).manual_seed(_seed)
389
+
390
+ pipeline_args = {
391
+ "prompt": prompt,
392
+ "negative_prompt": negative_prompt if negative_prompt else None,
393
+ "width": _width,
394
+ "height": _height,
395
+ "num_inference_steps": _num_inference_steps,
396
+ "guidance_scale": _guidance_scale,
397
+ "generator": generator,
398
+ "num_images_per_prompt": kwargs.get("num_images_per_prompt", 1)
399
+ }
400
+ if "eta" in kwargs: pipeline_args["eta"] = kwargs["eta"]
401
+ if "clip_skip" in kwargs and hasattr(self.pipeline, "clip_skip"): # Handle clip_skip if supported
402
+ pipeline_args["clip_skip"] = kwargs["clip_skip"]
403
+
404
+
405
+ ASCIIColors.info(f"Generating image with prompt: '{prompt[:100]}...'")
406
+ ASCIIColors.debug(f"Pipeline args: {pipeline_args}")
407
+
408
+ try:
409
+ with torch.no_grad(): # Important for inference
410
+ pipeline_output = self.pipeline(**pipeline_args)
411
+
412
+ pil_image: Image.Image = pipeline_output.images[0]
413
+
414
+ # Convert PIL Image to bytes (PNG)
415
+ img_byte_arr = BytesIO()
416
+ pil_image.save(img_byte_arr, format="PNG")
417
+ img_bytes = img_byte_arr.getvalue()
418
+
419
+ ASCIIColors.green("Image generated successfully.")
420
+ return img_bytes
421
+
422
+ except Exception as e:
423
+ trace_exception(e)
424
+ raise Exception(f"Diffusers image generation failed: {e}") from e
425
+
426
+ def list_services(self, **kwargs) -> List[Dict[str, str]]:
427
+ """
428
+ Lists the currently loaded model as the available service.
429
+ Future: Could scan local model directories or list known HF models.
430
+ """
431
+ if self.pipeline and self.current_model_id_or_path:
432
+ return [{
433
+ "name": os.path.basename(self.current_model_id_or_path),
434
+ "caption": f"Diffusers: {os.path.basename(self.current_model_id_or_path)}",
435
+ "help": (f"Currently loaded model. Path/ID: {self.current_model_id_or_path}. "
436
+ f"Device: {self.config['device']}. DType: {self.config['torch_dtype_str']}. "
437
+ f"Scheduler: {self.pipeline.scheduler.__class__.__name__}.")
438
+ }]
439
+ return [{"name": "diffusers_unloaded", "caption": "No Diffusers model loaded", "help": "Configure a model in settings."}]
440
+
441
+ def get_settings(self, **kwargs) -> List[Dict[str, Any]]:
442
+ """
443
+ Retrieves the current configurable settings for the Diffusers binding.
444
+ """
445
+ # Actual device and dtype after auto-resolution
446
+ resolved_device = self.config['device']
447
+ resolved_dtype_str = self.config['torch_dtype_str']
448
+
449
+ # For display, show the original 'auto' if it was set that way, plus the resolved value
450
+ display_device = self.config['device'] if self.config['device'].lower() != 'auto' else f"auto ({resolved_device})"
451
+ display_dtype = self.config['torch_dtype_str'] if self.config['torch_dtype_str'].lower() != 'auto' else f"auto ({resolved_dtype_str})"
452
+
453
+ settings = [
454
+ {"name": "model_id_or_path", "type": "str", "value": self.config["model_id_or_path"], "description": "Hugging Face model ID or local path to Diffusers model directory."},
455
+ {"name": "device", "type": "str", "value": self.config["device"], "description": f"Device for inference. Current resolved: {resolved_device}", "options": ["auto", "cuda", "mps", "cpu"]},
456
+ {"name": "torch_dtype_str", "type": "str", "value": self.config["torch_dtype_str"], "description": f"Torch dtype for model. Current resolved: {resolved_dtype_str}", "options": ["auto", "float16", "bfloat16", "float32"]},
457
+ {"name": "hf_variant", "type": "str", "value": self.config["hf_variant"], "description": "Model variant (e.g., 'fp16', 'bf16'). Optional."},
458
+ {"name": "use_safetensors", "type": "bool", "value": self.config["use_safetensors"], "description": "Prefer loading models from .safetensors files."},
459
+ {"name": "scheduler_name", "type": "str", "value": self.config["scheduler_name"], "description": "Scheduler to use for diffusion.", "options": list(SCHEDULER_MAPPING.keys())},
460
+ {"name": "safety_checker_on", "type": "bool", "value": self.config["safety_checker_on"], "description": "Enable the safety checker (if model has one)."},
461
+ {"name": "enable_cpu_offload", "type": "bool", "value": self.config["enable_cpu_offload"], "description": "Enable model CPU offload (saves VRAM, slower)."},
462
+ {"name": "enable_sequential_cpu_offload", "type": "bool", "value": self.config["enable_sequential_cpu_offload"], "description": "Enable sequential CPU offload (more VRAM savings, much slower)."},
463
+ {"name": "enable_xformers", "type": "bool", "value": self.config["enable_xformers"], "description": "Enable xFormers memory efficient attention (if available)."},
464
+ {"name": "default_width", "type": "int", "value": self.config["default_width"], "description": "Default width for generated images."},
465
+ {"name": "default_height", "type": "int", "value": self.config["default_height"], "description": "Default height for generated images."},
466
+ {"name": "num_inference_steps", "type": "int", "value": self.config["num_inference_steps"], "description": "Default number of inference steps."},
467
+ {"name": "guidance_scale", "type": "float", "value": self.config["guidance_scale"], "description": "Default guidance scale (CFG)."},
468
+ {"name": "seed", "type": "int", "value": self.config["seed"], "description": "Default seed for generation (-1 for random)."},
469
+ {"name": "hf_token", "type": "str", "value": self.config["hf_token"], "description": "Hugging Face API token (for private/gated models). Set to 'None' or empty if not needed. Store securely.", "is_secret": True},
470
+ {"name": "local_files_only", "type": "bool", "value": self.config["local_files_only"], "description": "Only use local files, do not try to download."},
471
+ ]
472
+ return settings
473
+
474
+ def set_settings(self, settings: Union[Dict[str, Any], List[Dict[str, Any]]], **kwargs) -> bool:
475
+ """
476
+ Applies new settings to the Diffusers binding. Some settings may trigger a model reload.
477
+ """
478
+ if isinstance(settings, list): # Convert from ConfigTemplate list format
479
+ parsed_settings = {item["name"]: item["value"] for item in settings if "name" in item and "value" in item}
480
+ elif isinstance(settings, dict):
481
+ parsed_settings = settings
482
+ else:
483
+ ASCIIColors.error("Invalid settings format. Expected a dictionary or list of dictionaries.")
484
+ return False
485
+
486
+ old_config = self.config.copy()
487
+ needs_reload = False
488
+
489
+ for key, value in parsed_settings.items():
490
+ if key in self.config:
491
+ if self.config[key] != value:
492
+ self.config[key] = value
493
+ ASCIIColors.info(f"Setting '{key}' changed to: {value}")
494
+ if key in ["model_id_or_path", "device", "torch_dtype_str",
495
+ "use_safetensors", "safety_checker_on", "hf_variant",
496
+ "enable_cpu_offload", "enable_sequential_cpu_offload", "enable_xformers",
497
+ "hf_token", "local_files_only"]:
498
+ needs_reload = True
499
+ elif key == "scheduler_name" and self.pipeline: # Scheduler can be changed on loaded pipeline
500
+ self._set_scheduler() # Attempt to apply immediately
501
+ else:
502
+ ASCIIColors.warning(f"Unknown setting '{key}' ignored.")
503
+
504
+ if needs_reload:
505
+ ASCIIColors.info("Reloading model due to settings changes...")
506
+ try:
507
+ # Resolve auto device/dtype again if they were part of the change
508
+ if "device" in parsed_settings and self.config["device"].lower() == "auto":
509
+ if torch.cuda.is_available(): self.config["device"] = "cuda"
510
+ elif torch.backends.mps.is_available(): self.config["device"] = "mps"
511
+ else: self.config["device"] = "cpu"
512
+
513
+ if "torch_dtype_str" in parsed_settings and self.config["torch_dtype_str"].lower() == "auto":
514
+ self.config["torch_dtype_str"] = "float16" if self.config["device"] != "cpu" else "float32"
515
+
516
+ # Update torch_dtype object from string
517
+ self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
518
+ if self.torch_dtype == "auto": # Should be resolved by now
519
+ self.torch_dtype = torch.float16 if self.config["device"] != "cpu" else torch.float32
520
+ self.config["torch_dtype_str"] = TORCH_DTYPE_MAP_OBJ_TO_STR.get(self.torch_dtype, "float32")
521
+
522
+
523
+ self.load_model()
524
+ ASCIIColors.green("Model reloaded successfully with new settings.")
525
+ except Exception as e:
526
+ trace_exception(e)
527
+ ASCIIColors.error(f"Failed to reload model with new settings: {e}. Reverting critical settings.")
528
+ # Revert critical settings and try to reload with old config
529
+ self.config = old_config
530
+ self.torch_dtype = TORCH_DTYPE_MAP_STR_TO_OBJ.get(self.config["torch_dtype_str"].lower(), torch.float32)
531
+ try:
532
+ self.load_model()
533
+ ASCIIColors.info("Reverted to previous model configuration.")
534
+ except Exception as e_revert:
535
+ trace_exception(e_revert)
536
+ ASCIIColors.error(f"Failed to revert to previous model configuration: {e_revert}. Binding may be unstable.")
537
+ return False
538
+ return True
539
+
540
+ def __del__(self):
541
+ self.unload_model()
542
+
543
+ # Example Usage (for testing within this file)
544
+ if __name__ == '__main__':
545
+ ASCIIColors.magenta("--- Diffusers TTI Binding Test ---")
546
+
547
+ if not DIFFUSERS_AVAILABLE:
548
+ ASCIIColors.error("Diffusers or its dependencies are not available. Cannot run test.")
549
+ # Attempt to guide user for installation
550
+ print("Please ensure PyTorch, Diffusers, Pillow, and Transformers are installed.")
551
+ print("For PyTorch with CUDA: visit https://pytorch.org/get-started/locally/")
552
+ print("Then: pip install diffusers Pillow transformers safetensors")
553
+ exit(1)
554
+
555
+ # --- Configuration ---
556
+ # Small, fast model for testing. Replace with a full model for real use.
557
+ # "CompVis/stable-diffusion-v1-4" is ~5GB
558
+ # "google/ddpm-cat-256" is smaller, but a DDPM, not Stable Diffusion.
559
+ # Using a tiny SD model if one exists, or a small variant.
560
+ # For a quick test, let's try a small LCM LoRA with SD1.5 if possible or just a base model.
561
+ # Note: "runwayml/stable-diffusion-v1-5" is a good standard test model.
562
+ # For a *very* quick CI-like test, one might use a dummy model or a very small one.
563
+ # Let's use a smaller SD variant if available, otherwise default to 2.1-base.
564
+ test_model_id = "runwayml/stable-diffusion-v1-5" # ~4GB download. Use a smaller one if you have it locally.
565
+ # test_model_id = "hf-internal-testing/tiny-stable-diffusion-pipe" # Very small, for testing structure
566
+
567
+ # Create dummy lollms_paths
568
+ temp_paths_dir = Path(__file__).parent / "temp_lollms_paths_diffusers"
569
+ temp_paths_dir.mkdir(parents=True, exist_ok=True)
570
+ mock_lollms_paths = {
571
+ "personal_models_path": temp_paths_dir / "personal_models",
572
+ "models_zoo_path": temp_paths_dir / "models_zoo",
573
+ "shared_cache_path": temp_paths_dir / "shared_cache", # For Hugging Face cache
574
+ }
575
+ for p in mock_lollms_paths.values(): Path(p).mkdir(parents=True, exist_ok=True)
576
+ (Path(mock_lollms_paths["personal_models_path"]) / "diffusers_models").mkdir(exist_ok=True)
577
+
578
+
579
+ binding_config = {
580
+ "model_id_or_path": test_model_id,
581
+ "device": "auto", # Let it auto-detect
582
+ "torch_dtype_str": "auto",
583
+ "num_inference_steps": 10, # Faster for testing
584
+ "default_width": 256, # Smaller for faster testing
585
+ "default_height": 256,
586
+ "safety_checker_on": False, # Often disabled for local use flexibility
587
+ "hf_variant": "fp16" if test_model_id == "runwayml/stable-diffusion-v1-5" else None, # SD 1.5 has fp16 variant
588
+ }
589
+
590
+ try:
591
+ ASCIIColors.cyan("\n1. Initializing DiffusersTTIBinding_Impl...")
592
+ binding = DiffusersTTIBinding_Impl(config=binding_config, lollms_paths=mock_lollms_paths)
593
+ ASCIIColors.green("Initialization successful.")
594
+ ASCIIColors.info(f"Loaded model: {binding.current_model_id_or_path}")
595
+ ASCIIColors.info(f"Device: {binding.config['device']}, DType: {binding.config['torch_dtype_str']}")
596
+ ASCIIColors.info(f"Scheduler: {binding.pipeline.scheduler.__class__.__name__ if binding.pipeline else 'N/A'}")
597
+
598
+
599
+ ASCIIColors.cyan("\n2. Listing services...")
600
+ services = binding.list_services()
601
+ ASCIIColors.info(json.dumps(services, indent=2))
602
+ assert services and services[0]["name"] == os.path.basename(binding.current_model_id_or_path)
603
+
604
+ ASCIIColors.cyan("\n3. Getting settings...")
605
+ settings_list = binding.get_settings()
606
+ ASCIIColors.info(json.dumps(settings_list, indent=2, default=str)) # default=str for Path objects if any
607
+ # Find model_id_or_path in settings
608
+ found_model_setting = any(s['name'] == 'model_id_or_path' and s['value'] == test_model_id for s in settings_list)
609
+ assert found_model_setting, "Model ID not found or incorrect in get_settings"
610
+
611
+
612
+ ASCIIColors.cyan("\n4. Generating an image...")
613
+ test_prompt = "A vibrant cat astronaut exploring a neon galaxy"
614
+ test_negative_prompt = "blurry, low quality, text, watermark"
615
+
616
+ # Use smaller dimensions for test if default are large
617
+ gen_width = min(binding.config["default_width"], 256)
618
+ gen_height = min(binding.config["default_height"], 256)
619
+
620
+ image_bytes = binding.generate_image(
621
+ prompt=test_prompt,
622
+ negative_prompt=test_negative_prompt,
623
+ width=gen_width, height=gen_height,
624
+ num_inference_steps=8 # Even fewer for speed
625
+ )
626
+ assert image_bytes and isinstance(image_bytes, bytes)
627
+ ASCIIColors.green(f"Image generated successfully (size: {len(image_bytes)} bytes).")
628
+ # Save the image for verification
629
+ test_image_path = Path(__file__).parent / "test_diffusers_image.png"
630
+ with open(test_image_path, "wb") as f:
631
+ f.write(image_bytes)
632
+ ASCIIColors.info(f"Test image saved to: {test_image_path.resolve()}")
633
+
634
+
635
+ ASCIIColors.cyan("\n5. Setting new settings (changing scheduler and guidance_scale)...")
636
+ new_settings_dict = {
637
+ "scheduler_name": "ddim", # Change scheduler
638
+ "guidance_scale": 5.0, # Change guidance scale
639
+ "num_inference_steps": 12 # Change inference steps
640
+ }
641
+ binding.set_settings(new_settings_dict)
642
+ assert binding.config["scheduler_name"] == "ddim"
643
+ assert binding.config["guidance_scale"] == 5.0
644
+ assert binding.config["num_inference_steps"] == 12
645
+ ASCIIColors.info(f"New scheduler (intended): ddim, Actual: {binding.pipeline.scheduler.__class__.__name__}")
646
+ ASCIIColors.info(f"New guidance_scale: {binding.config['guidance_scale']}")
647
+
648
+ ASCIIColors.cyan("\n6. Generating another image with new settings...")
649
+ image_bytes_2 = binding.generate_image(
650
+ prompt="A serene landscape with a crystal river",
651
+ width=gen_width, height=gen_height
652
+ )
653
+ assert image_bytes_2 and isinstance(image_bytes_2, bytes)
654
+ ASCIIColors.green(f"Second image generated successfully (size: {len(image_bytes_2)} bytes).")
655
+ test_image_path_2 = Path(__file__).parent / "test_diffusers_image_2.png"
656
+ with open(test_image_path_2, "wb") as f:
657
+ f.write(image_bytes_2)
658
+ ASCIIColors.info(f"Second test image saved to: {test_image_path_2.resolve()}")
659
+
660
+ # Test model reload by changing a critical parameter (e.g. safety_checker_on)
661
+ # This requires a different model or a config that can be easily toggled.
662
+ # For now, assume reload on critical param change works if no error is thrown.
663
+ ASCIIColors.cyan("\n7. Testing settings change requiring model reload (safety_checker_on)...")
664
+ current_safety_on = binding.config["safety_checker_on"]
665
+ binding.set_settings({"safety_checker_on": not current_safety_on})
666
+ assert binding.config["safety_checker_on"] == (not current_safety_on)
667
+ ASCIIColors.green("Model reload due to safety_checker_on change seems successful.")
668
+
669
+
670
+ except Exception as e:
671
+ trace_exception(e)
672
+ ASCIIColors.error(f"Diffusers binding test failed: {e}")
673
+ finally:
674
+ ASCIIColors.cyan("\nCleaning up...")
675
+ if 'binding' in locals() and binding:
676
+ binding.unload_model()
677
+
678
+ # Clean up temp_lollms_paths
679
+ import shutil
680
+ if temp_paths_dir.exists():
681
+ try:
682
+ shutil.rmtree(temp_paths_dir)
683
+ ASCIIColors.info(f"Cleaned up temporary directory: {temp_paths_dir}")
684
+ except Exception as e_clean:
685
+ ASCIIColors.warning(f"Could not fully clean up {temp_paths_dir}: {e_clean}")
686
+ if 'test_image_path' in locals() and test_image_path.exists():
687
+ # os.remove(test_image_path) # Keep for manual check
688
+ pass
689
+ if 'test_image_path_2' in locals() and test_image_path_2.exists():
690
+ # os.remove(test_image_path_2) # Keep for manual check
691
+ pass
692
+ ASCIIColors.magenta("--- Diffusers TTI Binding Test Finished ---")