File size: 1,609 Bytes
6c3f4f2 c68a9eb 6c3f4f2 687bf24 1d42f9c c68a9eb 6c3f4f2 687bf24 50936f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
# from diffusers import StableDiffusionPipeline
# def generate_image(prompt):
# model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
# model.to("cuda") # Use GPU for faster generation
# image = model(prompt).images[0]
# image.save("output.png")
# return "output.png"
# if __name__ == "__main__":
# prompt = "A friendly person saying 'How are you?'"
# print("Generated Image Path:", generate_image(prompt))
# from diffusers import StableDiffusionPipeline
# import torch
# def generate_image(prompt):
# model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
# # Use GPU if available, otherwise fallback to CPU
# device = "cuda" if torch.cuda.is_available() else "cpu"
# model.to(device)
# image = model(prompt).images[0]
# image.save("output.png")
# return "output.png"
# if __name__ == "__main__":
# prompt = "A friendly person saying 'How are you?'"
# print("Generated Image Path:", generate_image(prompt))
import spaces
from diffusers import StableDiffusionPipeline
import torch
# Preload the model globally
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1").to(device)
@spaces.GPU
def generate_image(prompt):
"""Generate an image based on the input prompt."""
with torch.no_grad():
image = pipeline(prompt).images[0]
# Save the image locally and return the file path
image_path = "generated_image.png"
image.save(image_path)
return image_path
|