Qwenの画像編集モデル「Qwen−Image-Edit」を試す

uv init -p 3.12 qwen-image-edit-work && cd $_

uv add dfloat11[cuda12]

 + dfloat11==0.3.2

uv add git+https://github.com/huggingface/diffusers

 + diffusers==0.36.0.dev0 (from git+https://github.com/huggingface/diffusers@4fcd0bc7ebb934a1559d0b516f09534ba22c8a0d)

uv add torchvision

import argparse
import torch
from diffusers.utils import load_image
from diffusers import QwenImageTransformer2DModel, QwenImageEditPipeline
from transformers.modeling_utils import no_init_weights
from dfloat11 import DFloat11Model

def parse_args():
    parser = argparse.ArgumentParser(description='Edit images using Qwen-Image-Edit model')
    parser.add_argument('--cpu_offload', action='store_true', help='Enable CPU offloading')
    parser.add_argument('--cpu_offload_blocks', type=int, default=16, help='Number of transformer blocks to offload to CPU')
    parser.add_argument('--no_pin_memory', action='store_true', help='Disable memory pinning')
    parser.add_argument('--image', type=str, default="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png",
                        help='Path to input image or URL')
    parser.add_argument('--prompt', type=str, default='Add a hat to the cat.',
                        help='Text prompt for image editing')
    parser.add_argument('--negative_prompt', type=str, default=' ',
                        help='Negative prompt for image editing')
    parser.add_argument('--num_inference_steps', type=int, default=50,
                        help='Number of denoising steps')
    parser.add_argument('--true_cfg_scale', type=float, default=4.0,
                        help='Classifier free guidance scale')
    parser.add_argument('--seed', type=int, default=42,
                        help='Random seed for generation')
    parser.add_argument('--output', type=str, default='qwen_image_edit.png',
                        help='Output image path')
    return parser.parse_args()

args = parse_args()
model_id = "Qwen/Qwen-Image-Edit"

with no_init_weights():
    transformer = QwenImageTransformer2DModel.from_config(
        QwenImageTransformer2DModel.load_config(
            model_id, subfolder="transformer",
        ),
    ).to(torch.bfloat16)

DFloat11Model.from_pretrained(
    "DFloat11/Qwen-Image-Edit-DF11",
    device="cpu",
    cpu_offload=args.cpu_offload,
    cpu_offload_blocks=args.cpu_offload_blocks,
    pin_memory=not args.no_pin_memory,
    bfloat16_model=transformer,
)

pipeline = QwenImageEditPipeline.from_pretrained(
    model_id, transformer=transformer, torch_dtype=torch.bfloat16,
)
pipeline.enable_model_cpu_offload()
pipeline.set_progress_bar_config(disable=None)

image = load_image(args.image)
inputs = {
    "image": image,
    "prompt": args.prompt,
    "generator": torch.manual_seed(args.seed),
    "true_cfg_scale": args.true_cfg_scale,
    "negative_prompt": args.negative_prompt,
    "num_inference_steps": args.num_inference_steps,
}

with torch.inference_mode():
    output = pipeline(**inputs)
    output_image = output.images[0]
    output_image.save(args.output)

max_gpu_memory = torch.cuda.max_memory_allocated()
print(f"Max GPU memory allocated: {max_gpu_memory / 1000 ** 3:.2f} GB")

uv run qwen_image_edit.py --help

usage: qwen_image_edit.py [-h] [--cpu_offload] [--cpu_offload_blocks CPU_OFFLOAD_BLOCKS] [--no_pin_memory]
                          [--image IMAGE] [--prompt PROMPT] [--negative_prompt NEGATIVE_PROMPT]
                          [--num_inference_steps NUM_INFERENCE_STEPS] [--true_cfg_scale TRUE_CFG_SCALE] [--seed SEED]
                          [--output OUTPUT]

Edit images using Qwen-Image-Edit model

options:
  -h, --help            show this help message and exit
  --cpu_offload         Enable CPU offloading
  --cpu_offload_blocks CPU_OFFLOAD_BLOCKS
                        Number of transformer blocks to offload to CPU
  --no_pin_memory       Disable memory pinning
  --image IMAGE         Path to input image or URL
  --prompt PROMPT       Text prompt for image editing
  --negative_prompt NEGATIVE_PROMPT
                        Negative prompt for image editing
  --num_inference_steps NUM_INFERENCE_STEPS
                        Number of denoising steps
  --true_cfg_scale TRUE_CFG_SCALE
                        Classifier free guidance scale
  --seed SEED           Random seed for generation
  --output OUTPUT       Output image path

uv run qwen_image_edit.py \
    --cpu_offload \
    --image horserace.png \
    --prompt "Change the track surface from dirt to turf."

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 18.00 MiB. GPU 0 has a total capacity of 23.49 GiB of which 8.88 MiB is free. Including non-PyTorch memory, this process has 23.25 GiB memory in use. Of the allocated memory 22.57 GiB is allocated by PyTorch, and 230.19 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

uv run qwen_image_edit.py \
    --cpu_offload \
    --cpu_offload_blocks 120 \
    --image horserace.png \
    --no_pin_memory \
    --prompt "Change the track surface from dirt to turf."