""" Generator for SGLang Diffusion ComfyUI integration. """ import logging import os import psutil from comfy import model_detection, model_management from comfy.utils import ( calculate_parameters, load_torch_file, state_dict_prefix_replace, unet_to_diffusers, ) logger = logging.getLogger(__name__) class _SuppressModelConfigWarning(logging.Filter): """Suppress the expected 'Could not read model config' error from SGLang when loading safetensors files.""" def filter(self, record): return "Could not read model config" not in record.getMessage() logging.getLogger("sglang.multimodal_gen.registry").addFilter(_SuppressModelConfigWarning()) try: from sglang.multimodal_gen import DiffGenerator except ImportError: logger.error( "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'" ) from ..executors import ( AnimaExecutor, FluxExecutor, NativeExecutor, QwenImageEditExecutor, QwenImageExecutor, ZImageExecutor, ) from .model_patcher import SGLDModelPatcher class SGLDiffusionGenerator: """Generator for SGLang Diffusion models in ComfyUI.""" def __init__(self): self.model_path = None self.generator = None self.executor = None self.last_options = None self.pipeline_class_dict = { "flux": "ComfyUIFluxPipeline", "lumina2": "ComfyUIZImagePipeline", # zimage "qwen_image": "ComfyUIQwenImagePipeline", "qwen_image_edit": "ComfyUIQwenImageEditPipeline", "anima": None, # Native mode: no SGLang pipeline, run directly via ComfyUI } self.executor_class_dict = { "flux": FluxExecutor, "lumina2": ZImageExecutor, "qwen_image": QwenImageExecutor, "qwen_image_edit": QwenImageEditExecutor, "anima": AnimaExecutor, } def __del__(self): self.close_generator() def init_generator( self, model_path: str, pipeline_class_name: str, kwargs: dict = None ): """Initialize the diffusion generator.""" if self.generator is not None: return self.generator if kwargs is None: kwargs = {} # Set comfyui_mode for ComfyUI integration kwargs["comfyui_mode"] = True self.generator = DiffGenerator.from_pretrained( model_path=model_path, pipeline_class_name=pipeline_class_name, **kwargs, ) return self.generator def kill_generator(self): """Kill worker processes manually because generator shutdown cannot terminate them.""" current_pid = os.getpid() worker_processes = [] for proc in psutil.process_iter(["pid", "name", "cmdline"]): try: # Look for sglang-diffusionWorker processes if proc.info["cmdline"]: cmdline = " ".join(proc.info["cmdline"]) if "sgl_diffusion::" in cmdline: if proc.info["pid"] != current_pid: worker_processes.append(proc) except (psutil.NoSuchProcess, psutil.AccessDenied): continue if worker_processes: logger.info( f"Found {len(worker_processes)} worker processes to terminate..." ) for proc in worker_processes: try: logger.info( f"Terminating worker process {proc.info['pid']}: {proc.info['name']}" ) proc.terminate() proc.wait(timeout=5) except psutil.TimeoutExpired: logger.warning( f"Process {proc.info['pid']} did not terminate, forcing kill..." ) try: proc.kill() proc.wait(timeout=2) except (psutil.NoSuchProcess, psutil.TimeoutExpired): pass except (psutil.NoSuchProcess, psutil.AccessDenied): pass def close_generator(self): """Close and cleanup the generator and all associated resources.""" if self.generator is not None: self.generator.shutdown() self.kill_generator() # Clear other references self.last_options = None self.model_path = None self.generator = None self.executor = None def get_comfyui_model(self, model_path: str, model_options: dict = None): """Get ComfyUI model from model path.""" if model_options is None: model_options = {} dtype = model_options.get("dtype", None) # Allow loading unets from checkpoint files sd = load_torch_file(model_path) diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) temp_sd = state_dict_prefix_replace( sd, {diffusion_model_prefix: ""}, filter_keys=True ) if len(temp_sd) > 0: sd = temp_sd parameters = calculate_parameters(sd) load_device = model_management.get_torch_device() model_detect_config = model_detection.detect_unet_config(sd, "") model_type = model_detect_config.get("image_model", None) # Native mode when: model type not in dict (unknown type) OR value is None (explicit native entry) is_native = self.pipeline_class_dict.get(model_type) is None if is_native: logger.info(f"Model type '{model_type}' has no SGLang pipeline, using native mode") model_config = model_detection.model_config_from_unet(sd, "") if model_config is not None: new_sd = sd else: new_sd = model_detection.convert_diffusers_mmdit(sd, "") if new_sd is not None: # diffusers mmdit model_config = model_detection.model_config_from_unet(new_sd, "") if model_config is None: return None else: # diffusers unet model_config = model_detection.model_config_from_diffusers_unet(sd) if model_config is None: return None diffusers_keys = unet_to_diffusers(model_config.unet_config) new_sd = {} for k in diffusers_keys: if k in sd: new_sd[diffusers_keys[k]] = sd.pop(k) offload_device = model_management.unet_offload_device() if dtype is None: unet_dtype = model_management.unet_dtype( model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes, ) else: unet_dtype = dtype manual_cast_dtype = model_management.unet_manual_cast( unet_dtype, load_device, model_config.supported_inference_dtypes ) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) model_config.custom_operations = model_options.get("custom_operations", None) if is_native: # Native mode: load actual model weights without SGLang comfyui_model = model_config.get_model(new_sd, "") comfyui_model = comfyui_model.to(offload_device) comfyui_model.load_model_weights(new_sd, "") else: model_config.unet_config["disable_unet_model_creation"] = True comfyui_model = model_config.get_model({}) return comfyui_model, model_config, model_type, is_native def load_model( self, model_path: str, model_options: dict = None, sgld_options: dict = None ): """Load model and return model patcher.""" gather_options = { "model_path": model_path, "model_options": model_options, "sgld_options": sgld_options, } if ( self.last_options is not None and self.last_options == gather_options and self.generator is not None ): return self.generator else: self.close_generator() self.last_options = gather_options self.model_path = model_path comfyui_model, model_config, model_type, is_native = self.get_comfyui_model( model_path, model_options ) set_model_type = sgld_options.pop("model_type", None) if sgld_options else None if set_model_type is not None and set_model_type in self.pipeline_class_dict: model_type = set_model_type is_native = self.pipeline_class_dict[model_type] is None if is_native: # Native mode: no SGLang, delegate forward() directly to ComfyUI model executor_class = self.executor_class_dict.get(model_type, NativeExecutor) native_diffusion_model = comfyui_model.diffusion_model self.executor = executor_class( None, model_path, comfyui_model, model_config, native_diffusion_model ) else: pipeline_class_name = self.pipeline_class_dict[model_type] executor_class = self.executor_class_dict[model_type] self.generator = self.init_generator( model_path, pipeline_class_name, sgld_options ) self.executor = executor_class( self.generator, model_path, comfyui_model, model_config ) comfyui_model.diffusion_model = self.executor load_device = model_management.get_torch_device() offload_device = model_management.unet_offload_device() return SGLDModelPatcher( comfyui_model, load_device, offload_device, model_type=model_type, native_mode=is_native )