feat: cp support in qwenimage 2.1

[sayakpaul]

Warning

The changes needed to facilitate CP in QwenImage 2.1 turned out to be more involved than in other models because of its KV caching machinery and how it disentangles the prefilling stage from decoding. The changes present in this PR, IMO, are minimal to accommodate CP in QwenImage 2.1. I am open to other alternatives, too.

Given the above information, one would ask why add CP support in the first place then?

TL;DR: CP enables Qwen-Image 2.1 multi-condition generation on two 24 GB GPUs when the same workload cannot run on one 24 GB GPU, reducing peak transformer-path memory from 26.52 GiB to 18.92 GiB per rank.

It is because when a user has, say, two consumer GPUs like A10G (each having 24 GB of VRAM), they can still perform inference with a reasonably large context window. Consider the below scenario:

Workload:

• Three distinct 1024² condition images
• 1024² output
• 40 steps, BF16, KV caching
• Qwen3-VL embeddings precomputed, then the text encoder removed

Script with HF Jobs

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "accelerate",
#   "diffusers @ git+https://github.com/huggingface/diffusers.git@cp-support-qwenimage2.1",
#   "pillow",
#   "protobuf",
#   "sentencepiece",
#   "torch",
#   "torchvision",
#   "transformers",
# ]
# ///

import argparse
import gc
import hashlib
import io
import json
import os
import subprocess
import sys
import time
from pathlib import Path

import torch
import torch.distributed as dist
from PIL import Image, ImageDraw

from diffusers import QwenImage21Pipeline
from diffusers.models._modeling_parallel import ContextParallelConfig


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cp-degree", type=int, default=1)
    parser.add_argument("--conditions", type=int, default=4)
    parser.add_argument("--resolution", type=int, default=1024)
    parser.add_argument("--steps", type=int, default=40)
    parser.add_argument("--worker", action="store_true")
    return parser.parse_args()


def make_images(count, resolution):
    images = []
    colors = [
        ((245, 225, 190), (78, 45, 25)),
        ((180, 220, 255), (35, 70, 130)),
        ((210, 245, 205), (40, 105, 50)),
        ((235, 205, 245), (95, 45, 115)),
    ]
    for index in range(count):
        background, foreground = colors[index % len(colors)]
        image = Image.new("RGB", (resolution, resolution), background)
        draw = ImageDraw.Draw(image)
        margin = resolution // 8
        offset = index * resolution // 32
        draw.ellipse(
            (margin + offset, margin, resolution - margin, resolution - margin - offset),
            fill=foreground,
        )
        draw.rectangle(
            (resolution // 3, resolution // 10 + offset, 2 * resolution // 3, resolution // 3),
            fill=(95, 55, 145),
        )
        images.append(image)
    return images


def gib(value):
    return value / 1024**3


def launch_workers(args):
    command = [
        sys.executable,
        "-m",
        "torch.distributed.run",
        "--standalone",
        f"--nproc_per_node={args.cp_degree}",
        os.path.abspath(__file__),
        "--worker",
        "--cp-degree",
        str(args.cp_degree),
        "--conditions",
        str(args.conditions),
        "--resolution",
        str(args.resolution),
        "--steps",
        str(args.steps),
    ]
    subprocess.run(command, check=True)


def main():
    args = parse_args()
    if args.cp_degree > 1 and not args.worker:
        launch_workers(args)
        return

    distributed = args.cp_degree > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    if distributed:
        dist.init_process_group("nccl", device_id=device)
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        rank = 0
        world_size = 1

    process_started_at = time.perf_counter()
    try:
        pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", dtype=torch.bfloat16)
        images = [image.convert("RGBA") for image in make_images(args.conditions, args.resolution)]
        input_images = [
            pipe.image_processor.resize(image, width=args.resolution, height=args.resolution) for image in images
        ]
        prompt = (
            "Combine the subject, setting, composition, and colors from all references into one coherent "
            "oil painting"
        )

        pipe.text_encoder.to(device)
        torch.cuda.reset_peak_memory_stats(device)
        encode_started_at = time.perf_counter()
        with torch.no_grad():
            prompt_embeds, prompt_embeds_mask, image_pad_mask = pipe.encode_prompt(
                prompt=prompt,
                image=input_images,
                device=device,
                num_images_per_prompt=1,
            )
        torch.cuda.synchronize(device)
        encode_seconds = time.perf_counter() - encode_started_at
        encoder_peak_gib = gib(torch.cuda.max_memory_allocated(device))

        pipe.text_encoder.to("cpu")
        pipe.text_encoder = None
        pipe.processor = None
        gc.collect()
        torch.cuda.empty_cache()

        pipe.to(device)
        if distributed:
            pipe.transformer.enable_parallelism(
                config=ContextParallelConfig(ulysses_degree=world_size, ulysses_anything=True),
            )
        pipe.encode_prompt = lambda **kwargs: (prompt_embeds, prompt_embeds_mask, image_pad_mask)
        pipe.set_progress_bar_config(disable=True)
        loaded_gib = gib(torch.cuda.memory_allocated(device))

        generator = torch.Generator(device=device).manual_seed(42)
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats(device)
        if distributed:
            dist.barrier()
        inference_started_at = time.perf_counter()

        def offload_before_decode(pipeline, step_index, timestep, callback_kwargs):
            if distributed and step_index == args.steps - 1:
                pipeline.transformer.to("cpu")
                gc.collect()
                torch.cuda.empty_cache()
            return callback_kwargs

        image = pipe(
            prompt_embeds=prompt_embeds,
            prompt_embeds_mask=prompt_embeds_mask,
            image=images,
            generator=generator,
            num_inference_steps=args.steps,
            output_resolution=args.resolution,
            use_kv_cache=True,
            callback_on_step_end=offload_before_decode if distributed else None,
        ).images[0]
        torch.cuda.synchronize(device)
        inference_seconds = time.perf_counter() - inference_started_at

        metrics = torch.tensor(
            [encode_seconds, encoder_peak_gib, inference_seconds, gib(torch.cuda.max_memory_allocated(device))],
            device=device,
            dtype=torch.float64,
        )
        if distributed:
            dist.all_reduce(metrics, op=dist.ReduceOp.MAX)

        if rank == 0:
            image_bytes = io.BytesIO()
            image.save(image_bytes, format="PNG")
            result = {
                "status": "ok",
                "torch_version": torch.__version__,
                "gpu": torch.cuda.get_device_name(device),
                "gpu_total_gib": gib(torch.cuda.get_device_properties(device).total_memory),
                "world_size": world_size,
                "conditions": args.conditions,
                "resolution": args.resolution,
                "steps": args.steps,
                "encoder_peak_allocated_gib": metrics[1].item(),
                "loaded_without_text_encoder_gib": loaded_gib,
                "inference_peak_allocated_gib_per_gpu": metrics[3].item(),
                "encoder_seconds": metrics[0].item(),
                "inference_seconds": metrics[2].item(),
                "process_seconds": time.perf_counter() - process_started_at,
                "output_size": list(image.size),
                "output_sha256": hashlib.sha256(image_bytes.getvalue()).hexdigest(),
            }
            artifact_stem = f"qwen21_cp{world_size}_c{args.conditions}_r{args.resolution}_s{args.steps}"
            artifact_dir = Path(__file__).resolve().parent
            image_path = artifact_dir / f"{artifact_stem}.png"
            result_path = artifact_dir / f"{artifact_stem}.json"
            image.save(image_path)
            result["output_artifact"] = str(image_path)
            result["metrics_artifact"] = str(result_path)
            result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
            print("BENCHMARK_RESULT=" + json.dumps(result, sort_keys=True), flush=True)
    finally:
        if distributed:
            dist.destroy_process_group()


if __name__ == "__main__":
    main()

Result:

3 input condition images:

(Edit prompt: Combine the subject, setting, composition, and colors from all references into one coherent oil painting)

[sayakpaul]

Cc: @naykun

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.