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)