import torch
from diffusers import ZImagePipeline
import faulthandler
import signal

faulthandler.register(signal.SIGUSR1)
taille=512
taille=taille - (taille % 16)
steps=128
# 1. Load the pipeline
# Use bfloat16 for optimal performance on supported GPUs
device = "xpu"
print(torch.__version__)
print(torch.xpu.is_available())

print(torch.xpu.device_count())
print(torch.xpu.get_device_name(0))


pipe = ZImagePipeline.from_pretrained(
    "Images",
    dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)

pipe.to("xpu")

print("Letsgo")

# [Optional] Attention Backend
# Diffusers uses SDPA by default. Switch to Flash Attention for better efficiency if supported:
# pipe.transformer.set_attention_backend("flash")    # Enable Flash-Attention-2
# pipe.transformer.set_attention_backend("_flash_3") # Enable Flash-Attention-3

# [Optional] Model Compilation
# Compiling the DiT model accelerates inference, but the first run will take longer to compile.
# pipe.transformer.compile()

# [Optional] CPU Offloading
# Enable CPU offloading for memory-constrained devices.
# pipe.enable_model_cpu_offload()


prompt = "Deux grenouilles vertes en train de faire du judo sur un tapis au milieu d'une forêt."

#prompt = "Deux grenouilles vertes en train de faire du judo sur un tapis de mousse posée sur de l'herbe au milieu d'une forêt, la première en kimono rouge tient la seconde en kimono bleu suspendue à l'air, un scorpion noir regarde la scène.. La scène est vue de coté"

#prompt = "Deux grenouilles vertes en train de faire du judo sur un tapis carré de mousse posée sur de l'herbe au milieu d'une forêt, la première en kimono rouge tient en la regardant la seconde en kimono bleu suspendue à l'air, celle ci regarde par terre. Elles sont tetes nues, un scorpion noir regarde la scène à coté d'une coccinelle . La scène est vue de coté"

#prompt = "Deux grenouilles vertes en train de faire du judo sur un tapis carré de mousse posée sur de l'herbe au milieu d'une forêt, la première en kimono rouge tient en la regardant la seconde en kimono bleu suspendue à l'air, celle ci regarde par terre et a la bouche ouverte. Elles sont tetes nues, un scorpion noir regarde la scène à coté d'une coccinelle rouge avec des points noirs . La scène est vue de coté"

# Là
#prompt = "Deux grenouilles vertes en train de faire du judo sur un tapis carré de mousse posée sur de l'herbe au milieu d'une forêt, la première en kimono rouge tient en la regardant la seconde en kimono bleu, suspendue en l'air,  qui regarde par terre et a la bouche ouverte. Elles sont tetes nues, un scorpion noir regarde la scène à coté d'une coccinelle rouge avec des points noirs . La scène est vue de coté"

#prompt = "Deux lapins en train de faire du judo sur un tapis carré de mousse posée sur de l'herbe au milieu d'une forêt, la première en kimono rouge tient en la regardant la seconde en kimono bleu suspendue à l'air, celle ci regarde par terre et a la bouche ouverte. Elles sont tetes nues, un scorpion noir regarde la scène à coté d'une coccinelle rouge avec des points noirs . La scène est vue de coté"

#prompt="Dans une forêt, il y a un tapis de mousse sur lequel deux lapins font du judo. L'un en kimono rouge tient en le regardant l'autre suspendu en l'air. Celui qui est suspendu en l'air a la boucge ouverte et regarde par terre. Un scorpion noir et une coccinelle regardent. La scène est vue de coté"

#prompt="Dans une forêt, il y a un tapis de mousse sur lequel deux lapin font du judo. L'un en kimono rouge tient en le regardant l'autre suspendu en l'air. Celui qui est suspendu en l'air a la boucge ouverte et regarde par terre. Un scorpion noir et une coccinelle regardent. La scène est vue de coté"

#prompt = "image précédente"

#prompt = "Deux hommes en train de faire du judo sur un tapis carré de mousse posée sur de l'herbe au milieu d'une forêt, le premier en kimono rouge tient en le regardant le second en kimono bleu suspendu en l'air, il regarde par terre et a la bouche ouverte. Ils sont tetes nues, un scorpion noir regarde la scène à coté d'une coccinelle rouge avec des points noirs . La scène est vue de coté"

#prompt="Femme sur une plage"

#prompt=""

# 2. Generate Image
image = pipe(
    prompt=prompt,
    height=taille,
    width=taille,
    num_inference_steps=steps,  # This actually results in 8 DiT forwards, était à 9
    guidance_scale=0.0,     # Guidance should be 0 for the Turbo models
    generator=torch.Generator("xpu").manual_seed(42),   # était 42, puis 12
).images[0]

#torch.xpu.empty_cache()

image.save("example.png")