Implementation:Microsoft DeepSpeedExamples Local Stable Diffusion Pipeline

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Generative AI, Computer Vision, Diffusion Models
Last Updated	2026-02-07 12:00 GMT

Overview

A local implementation of the Stable Diffusion text-to-image generation pipeline, adapted from HuggingFace Diffusers for use within the DeepSpeed training framework.

Description

This module provides a self-contained StableDiffusionPipeline class that inherits from DiffusionPipeline and implements the full text-to-image generation workflow. The pipeline orchestrates multiple components: a VAE (AutoencoderKL) for encoding/decoding images to/from latent space, a CLIP text encoder for converting text prompts to embeddings, a UNet2DConditionModel for iterative denoising of latent representations, and a configurable noise scheduler from the Karras diffusion family.

The __call__ method implements the complete inference pipeline: encoding text prompts (with optional classifier-free guidance via negative prompts), preparing initial latent noise, running the iterative denoising loop with the UNet, decoding latents back to pixel space through the VAE, and optionally running a safety checker to filter potentially harmful content. It supports advanced features including VAE slicing and tiling for memory-efficient processing of large images, sequential CPU offloading for running on limited GPU memory, and model CPU offloading for optimal memory management.

This local copy exists within the training directory to support distillation and LoRA training workflows where the pipeline is needed for validation image generation or teacher model inference during training. Key methods include _encode_prompt for text embedding with do_classifier_free_guidance support, prepare_latents for initializing noise tensors, and run_safety_checker for content filtering.

Usage

Use this pipeline for text-to-image generation within the Stable Diffusion training examples, particularly for validation during LoRA distillation training. It is used internally by the training scripts to generate sample images for quality assessment during training runs.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/stable_diffusion/local_pipeline_stable_diffusion.py
• Lines: 1-705

Signature

class StableDiffusionPipeline(DiffusionPipeline):
    def __init__(
        self,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        unet: UNet2DConditionModel,
        scheduler: KarrasDiffusionSchedulers,
        safety_checker: StableDiffusionSafetyChecker,
        feature_extractor: CLIPFeatureExtractor,
        requires_safety_checker: bool = True,
    ):
        ...

    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        callback: Optional[Callable] = None,
        callback_steps: int = 1,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
    ) -> StableDiffusionPipelineOutput:
        ...

Import

from local_pipeline_stable_diffusion import StableDiffusionPipeline

I/O Contract

Inputs

Name	Type	Required	Description
prompt	str or List[str]	No	Text prompt(s) to guide image generation
height	int	No	Height in pixels of generated image (default: unet sample_size * vae_scale_factor)
width	int	No	Width in pixels of generated image (default: unet sample_size * vae_scale_factor)
num_inference_steps	int	No	Number of denoising steps (default: 50)
guidance_scale	float	No	Classifier-free guidance scale (default: 7.5)
negative_prompt	str or List[str]	No	Prompt(s) for negative guidance
num_images_per_prompt	int	No	Number of images to generate per prompt (default: 1)
generator	torch.Generator	No	Random number generator for reproducibility
latents	torch.FloatTensor	No	Pre-generated noisy latents for generation

Outputs

Name	Type	Description
images	List[PIL.Image] or np.ndarray	Generated images in PIL format or as numpy arrays
nsfw_content_detected	List[bool]	Flags indicating whether safety checker detected NSFW content

Usage Examples

import torch
from local_pipeline_stable_diffusion import StableDiffusionPipeline

# Load pipeline from pretrained checkpoint
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

# Generate an image
prompt = "a photo of an astronaut riding a horse on mars"
result = pipe(prompt, num_inference_steps=50, guidance_scale=7.5)
image = result.images[0]
image.save("output.png")

# Generate with negative prompt
result = pipe(
    prompt="a beautiful landscape",
    negative_prompt="blurry, low quality",
    num_inference_steps=50
)