import os
from datetime import datetime

from gradio_client import Client


def save_images(images, used_seeds, output_dir="outputs"):
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")

    for index, image in enumerate(images):
        filename = f"{timestamp}_img{index + 1}.png"
        path = os.path.join(output_dir, filename)

        # `image` may be a PIL.Image or a path-like object depending on Gradio.
        if hasattr(image, "save"):
            image.save(path)
        else:
            from PIL import Image

            Image.open(image).save(path)

        print(f"Saved: {path}")

    print(f"Seeds used: {used_seeds}")


def main():
    # Space that runs the Z-Image UI; change if needed.
    space_id = "yingzhac/Z_image_NSFW"

    client = Client(space_id)

    prompt = "A beautiful portrait photo of a woman, 4k, highly detailed"
    negative_prompt = ""
    height = 1024
    width = 1024
    num_inference_steps = 9
    guidance_scale = 0.0  # 0 = no CFG, as recommended for Turbo
    seed = 42
    randomize_seed = False

    images, used_seeds = client.predict(
        prompt,
        negative_prompt,
        height,
        width,
        num_inference_steps,
        guidance_scale,
        seed,
        randomize_seed,
        api_name="/generate_image",
    )

    save_images(images, used_seeds)


if __name__ == "__main__":
    main()