CaesarAIART/caesaraiart.py

import torch
from diffusers import StableDiffusionPipeline

class CaesarAIART:
    def __init__(self,CURRENT_DIR=""):
        self.model_id = "CompVis/stable-diffusion-v1-4"
        self.device = "cuda"
        self.pipe = StableDiffusionPipeline.from_pretrained(self.model_id, torch_dtype=torch.float16)
        self.CURRENT_DIR = CURRENT_DIR
        self.pipe = self.pipe.to(self.device)
    def generate(self,prompt):
        image = self.pipe(prompt).images[0]  
        image.save(f"{self.CURRENT_DIR}/CaesarAIART/caesarart.png")

if __name__ == "__main__":
    prompt = "a photo of an astronaut riding a horse on mars"
    caesaraiart = CaesarAIART()
    caesaraiart.generate()