from diffusers import DiffusionPipeline
import torch

# Replace with your model path if it's not from the Hugging Face model hub
# model_id = "https://api-inference.huggingface.co/models/awabxnero/imagen"  # Your Hugging Face model repo
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")

pipe = pipe.to("cuda")  # Move model to GPU if available

# Function to generate image from prompt
def generate_image(prompt):
    # Generate image
    image = pipe(prompt).images[0]
    image.show()  # Show the image (you can also save or return the image)
    return image

if __name__ == "__main__":
    prompt = "A beautiful sunset over the ocean"
    generate_image(prompt)