imagen / model.py
awabxnero's picture
Update model.py
06789ee verified
raw
history blame
694 Bytes
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)