From 05f8a068f4683fc7c49d70d361f8765edf52617b Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Tue, 18 Nov 2025 23:30:57 +0530 Subject: [PATCH 1/2] Added Single File Diffusion models --- .../image_diffusion/diffusion_worker.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/api/transformerlab/plugins/image_diffusion/diffusion_worker.py b/api/transformerlab/plugins/image_diffusion/diffusion_worker.py index bfa7701b1..5699ce728 100644 --- a/api/transformerlab/plugins/image_diffusion/diffusion_worker.py +++ b/api/transformerlab/plugins/image_diffusion/diffusion_worker.py @@ -39,6 +39,21 @@ StableDiffusionXLControlNetImg2ImgPipeline, StableDiffusionXLControlNetUnionImg2ImgPipeline, StableDiffusionXLControlNetPAGImg2ImgPipeline, + StableDiffusionPipeline, + StableDiffusionImg2ImgPipeline, + StableDiffusionInpaintPipeline, + StableDiffusionXLImg2ImgPipeline, + StableDiffusionXLInpaintPipeline, + StableDiffusionXLInstructPix2PixPipeline, + StableDiffusionXLKDiffusionPipeline, + StableDiffusion3Pipeline, + LatentConsistencyModelPipeline, + LatentConsistencyModelImg2ImgPipeline, + StableDiffusionControlNetXSPipeline, + StableDiffusionXLControlNetXSPipeline, + LEditsPPPipelineStableDiffusion, + LEditsPPPipelineStableDiffusionXL, + PIAPipeline, FluxControlNetImg2ImgPipeline, StableDiffusionControlNetInpaintPipeline, StableDiffusionXLControlNetInpaintPipeline, @@ -70,6 +85,24 @@ "DPMSolverMultistepScheduler": DPMSolverMultistepScheduler, } +SINGLE_FILE_MAP = { + "StableDiffusionPipeline": StableDiffusionPipeline, + "StableDiffusionImg2ImgPipeline": StableDiffusionImg2ImgPipeline, + "StableDiffusionInpaintPipeline": StableDiffusionInpaintPipeline, + "StableDiffusionXLImg2ImgPipeline": StableDiffusionXLImg2ImgPipeline, + "StableDiffusionXLInpaintPipeline": StableDiffusionXLInpaintPipeline, + "StableDiffusionXLInstructPix2PixPipeline": StableDiffusionXLInstructPix2PixPipeline, + "StableDiffusionXLKDiffusionPipeline": StableDiffusionXLKDiffusionPipeline, + "StableDiffusion3Pipeline": StableDiffusion3Pipeline, + "LatentConsistencyModelPipeline": LatentConsistencyModelPipeline, + "LatentConsistencyModelImg2ImgPipeline": LatentConsistencyModelImg2ImgPipeline, + "StableDiffusionControlNetXSPipeline": StableDiffusionControlNetXSPipeline, + "StableDiffusionXLControlNetXSPipeline": StableDiffusionXLControlNetXSPipeline, + "LEditsPPPipelineStableDiffusion": LEditsPPPipelineStableDiffusion, + "LEditsPPPipelineStableDiffusionXL": LEditsPPPipelineStableDiffusionXL, + "PIAPipeline": PIAPipeline, +} + def load_controlnet_model(controlnet_id: str, device: str = "cuda") -> ControlNetModel: controlnet_model = ControlNetModel.from_pretrained( @@ -96,6 +129,13 @@ def latents_to_rgb(latents): return Image.fromarray(image_array) +def is_single_file_model(model_path): + """Check if the model is a single-file format (.safetensors or .ckpt)""" + if isinstance(model_path, str): + return model_path.endswith((".safetensors", ".ckpt", ".pt")) + return False + + def create_decode_callback(output_dir): """Create a callback function to decode and save latents at each step""" From 1541de657dd1f5e699f83ae8ff62f64770fc1e28 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Thu, 20 Nov 2025 00:53:30 +0530 Subject: [PATCH 2/2] Updates --- .../image_diffusion/diffusion_worker.py | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/api/transformerlab/plugins/image_diffusion/diffusion_worker.py b/api/transformerlab/plugins/image_diffusion/diffusion_worker.py index 5699ce728..0241d7586 100644 --- a/api/transformerlab/plugins/image_diffusion/diffusion_worker.py +++ b/api/transformerlab/plugins/image_diffusion/diffusion_worker.py @@ -28,6 +28,8 @@ EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler, ControlNetModel, + DDIMScheduler, + DDPMScheduler, StableDiffusionControlNetPAGPipeline, StableDiffusionXLControlNetPAGPipeline, FluxControlNetPipeline, @@ -83,6 +85,8 @@ "LMSDiscreteScheduler": LMSDiscreteScheduler, "EulerAncestralDiscreteScheduler": EulerAncestralDiscreteScheduler, "DPMSolverMultistepScheduler": DPMSolverMultistepScheduler, + "DDIMScheduler": DDIMScheduler, + "DDPMScheduler": DDPMScheduler, } SINGLE_FILE_MAP = { @@ -414,7 +418,6 @@ def load_pipeline_with_sharding( # Handle 'auto' device safely by falling back to cuda:0 for controlnet safe_device = "cuda:0" if device == "auto" else device - controlnet = controlnet_class.from_pretrained( controlnet_id, torch_dtype=torch.float16 if safe_device != "cpu" else torch.float32, @@ -638,6 +641,28 @@ def load_pipeline_with_device_map( torch.cuda.empty_cache() # Load appropriate pipeline + is_single_file = is_single_file_model(model_path) + if is_single_file: + if is_inpainting: + architecture = "StableDiffusionInpaintPipeline" + elif is_img2img: + architecture = "StableDiffusionImg2ImgPipeline" + else: + architecture = "StableDiffusionPipeline" + single_pipeline_class = SINGLE_FILE_MAP.get(architecture) + if not single_pipeline_class: + raise ValueError(f"Model architecture '{architecture}' not supported for single-file models") + + if single_pipeline_class: + pipe = single_pipeline_class.from_single_file( + model_path, + torch_dtype=torch.float16 if device != "cpu" else torch.float32, + ) + if device != "auto": + pipe = pipe.to(device) + return pipe + else: + print(f"Warning: No single-file pipeline class found for architecture '{architecture}'") if is_controlnet: CONTROLNET_PIPELINE_MAP = { "StableDiffusionPipeline": StableDiffusionControlNetPipeline,