infer_inpaint.py

import os
import time
from pathlib import Path

import torch
from diffusers import FlowMatchEulerDiscreteScheduler, GGUFQuantizationConfig
from diffusers.utils import load_image
from diffusers_local import patch # Apply necessary patches for local diffusers components

# 1. Import all necessary components
from diffusers_local.pipeline_z_image_control_unified import ZImageControlUnifiedPipeline
from diffusers_local.z_image_control_transformer_2d import ZImageControlTransformer2DModel

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,garbage_collection_threshold:0.7,max_split_size_mb:1024"

def main():
    # 1. Set params
    BASE_MODEL_ID = "."        
    GGUF_MODEL_FILE = "./transformer/z_image_turbo_control_unified_v2.1_q4_k_m.gguf"
    GGUF_MODEL_FILE = "./transformer/z_image_turbo_control_unified_v2.1_q8_0.gguf"
    
    use_gguf = True

    # prompt="一位年轻女子站在阳光明媚的海岸线上，白裙在轻拂的海风中微微飘动，裙摆轻盈飞扬。她拥有一头鲜艳的紫色长发，在风中轻盈舞动，发间系着一个精致的黑色蝴蝶结，与身后柔和的蔚蓝天空形成鲜明对比。她面容清秀，眉目精致，肤色白皙细腻，透着一股甜美的青春气息；神情柔和，略带羞涩，目光静静地凝望着远方的地平线，双手自然交叠于身前，手指清晰可见、五指完整、指节自然、姿势优雅放松，仿佛沉浸在思绪之中。背景是辽阔无垠、波光粼粼的大海，阳光洒在海面上，映出温暖的金色光晕，海浪轻轻拍打沙滩，天空湛蓝云朵稀薄。整体画面高清锐利、细节丰富、色彩鲜艳、焦点清晰、8K分辨率、杰作、最佳质量、无模糊、无噪点、无畸变、自然光照、电影级渲染。"
    prompt = "Photorealistic portrait of a beautiful young East Asian woman with long, vibrant purple hair and a black bow. She is wearing a flowing white summer dress, standing on a sunny beach with a sparkling ocean and clear blue sky in the background. Bright natural sunlight, sharp focus, ultra-detailed."
    negative_prompt = "Low quality, blurry, ugly, deformed fingers, extra fingers, bad hand, bad anatomy, noise, overexposed, underexposed"

    target_height = 1728
    target_width = 992
    num_inference_steps = 25
    guidance_scale = 0  # 2.5
    controlnet_conditioning_scale = 0.7
    controlnet_conditioning_refiner_scale = 0.75
    mask_blur_radius = 12
    seed = 48
    shift = 3.0
    inpaint_mode = "diff+inpaint"  # ("default", "diff", "diff+inpaint")
    generator = torch.Generator("cuda").manual_seed(seed)

    print("Loading Pipeline...")
    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)

    if use_gguf:
        transformer = ZImageControlTransformer2DModel.from_single_file(
            GGUF_MODEL_FILE,
            torch_dtype=torch.bfloat16,
            config=str(Path(GGUF_MODEL_FILE).parent),
            quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
            add_control_noise_refiner=True,  # <== If you don't want to use the control noise refiner disable here.
        )
    else:
        transformer = ZImageControlTransformer2DModel.from_pretrained(
            BASE_MODEL_ID,
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
            add_control_noise_refiner=True,  # <== If you don't want to use the control noise refiner disable here.
        )
    pipe = ZImageControlUnifiedPipeline.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.bfloat16,
        transformer=transformer,  # You don't need to load the transformer here if you don't intend to disable add_control_noise_refiner (only for diffusers model, gguf is required).
    )
    pipe.scheduler = scheduler

    # Apply optimization (Optional)
    pipe.enable_group_offload(
        onload_device="cuda", offload_device="cpu", offload_type="block_level", num_blocks_per_group=1, low_cpu_mem_usage=True, use_stream=True
    )
    pipe.vae.use_tiling = True
    # ---

    print("\nRunning Inference...")

    pose_image = load_image("assets/pose.jpg")
    inpaint_image = load_image("assets/inpaint.jpg")
    mask_image = load_image("assets/inpaint_mask.jpg")
    start_inference_time = time.time()

    generated_image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=inpaint_image,
        control_image=pose_image,
        mask_image=mask_image,
        mask_blur_radius=mask_blur_radius,        
        inpaint_mode=inpaint_mode,
        height=target_height,
        width=target_width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        controlnet_refiner_conditioning_scale=controlnet_conditioning_refiner_scale,
        generator=generator,
    ).images[0]

    end_inference_time = time.time()
    print(f"\nGeneration finished in {end_inference_time - start_inference_time:.2f} seconds.")

    # Save Output
    if not os.path.exists("outputs"):
        os.makedirs("outputs")
    output_filename = "outputs/z_image_controlnet_result_inpaint.png"
    generated_image.save(output_filename)
    print(f"Image successfully saved as '{output_filename}'")
    generated_image.show()


if __name__ == "__main__":
    main()