#!/usr/bin/env python
  """Test script for ZImage inpaint support."""

  import sys
  sys.path.insert(0, '/home/ohiom/diffusers/src')

  import torch
  import numpy as np
  from PIL import Image
  from diffusers import ZImageInpaintPipeline

  # Paths
  MODEL_PATH =
  "database/models/huggingface/models--Tongyi-MAI--Z-Image-Turbo/snapshots/78771b7e11b922c868dd766476bda1f4fc6bfc96"
  INPUT_IMAGE_PATH = "death_remix_1024.png"

  print("Loading ZImageInpaintPipeline...")
  pipe = ZImageInpaintPipeline.from_pretrained(
      MODEL_PATH,
      torch_dtype=torch.bfloat16,
      local_files_only=True,
  )
  pipe.to("cuda")
  print("Pipeline loaded.")

  # Load input image
  print(f"\nLoading input image from {INPUT_IMAGE_PATH}...")
  input_image = Image.open(INPUT_IMAGE_PATH).convert("RGB")
  print(f"Input image size: {input_image.size}")

  # Create a mask (white = inpaint, black = preserve)
  width, height = input_image.size
  mask = np.zeros((height, width), dtype=np.uint8)
  h_start, h_end = height // 4, 3 * height // 4
  w_start, w_end = width // 4, 3 * width // 4
  mask[h_start:h_end, w_start:w_end] = 255
  mask_image = Image.fromarray(mask)

  # Generate an inpainted image
  prompt = "a woman with pale skin in a black shirt, oil painting style"
  strength = 0.75

  print(f"\nGenerating inpainted image with prompt: {prompt}")
  print(f"Strength: {strength}")

  image = pipe(
      prompt=prompt,
      image=input_image,
      mask_image=mask_image,
      strength=strength,
      num_inference_steps=8,
      guidance_scale=1.0,
      generator=torch.Generator(device="cuda").manual_seed(42),
  ).images[0]

  output_path = "test_zimage_inpaint_output.png"
  image.save(output_path)
  print(f"\nImage saved to {output_path}")