from diffusers import DiffusionPipeline
import torch
import matplotlib.pyplot as plt

# Vérifier si un GPU est disponible et l'utiliser si possible
device = "cuda" if torch.cuda.is_available() else "cpu"

# Charger le modèle
pipe = DiffusionPipeline.from_pretrained("XLabs-AI/flux-controlnet-canny").to(device)

# Définir le prompt
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

# Générer l'image
image = pipe(prompt).images[0]

# Afficher l'image
plt.imshow(image)
plt.axis("off")
plt.show()