diff --git a/main.py b/main.py new file mode 100644 index 0000000..1b7b1ad --- /dev/null +++ b/main.py @@ -0,0 +1,30 @@ +import torch +from diffusers import StableDiffusionPipeline +from dotenv import load_dotenv +from os import getenv +from PIL.PngImagePlugin import PngInfo +import uuid +from random import randint + +load_dotenv() + +pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=getenv("TOKEN")) +pipe.to("cuda") + + +prompt = "red horse jumping over a dog" + +seed = randint(0, 2**64) +generator = torch.Generator("cuda").manual_seed(seed) +resp = pipe(prompt, generator=generator) +print(resp) +image = resp.images[0] + +metadata = PngInfo() +metadata.add_text("text", prompt) +metadata.add_text("seed", str(seed)) + +id = str(uuid.uuid4()) +print(id) + +image.save(f'/images/{id}.png', pnginfo=metadata) \ No newline at end of file