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) | |