# conda activate PyTorch
import torch

import os

from os.path import join, dirname

from diffusers import StableDiffusionPipeline

from dotenv import load_dotenv

load_dotenv(verbose=True)

dotenv_path = join(dirname(__file__), '.env')

load_dotenv(dotenv_path)

AUTH_TOKEN = os.environ.get("AUTH_TOKEN")

pipe = StableDiffusionPipeline.from_pretrained(

    "CompVis/stable-diffusion-v1-4",

    revision="fp16",

    torch_dtype=torch.float16, use_auth_token=AUTH_TOKEN)
pipe.to("cuda")

prompt = "concept art. Illustration,hyper quality,highly detailed, cinematic lighting. A teenage Japanese youth with neko cat ears on head top and neko cat tail on waist, He has a beautiful face with a hint of childhood. equipped with post-apocalyptic armour and tactical hand axe, is looking at us against an abandoned Tokyo building. He is wearing jeans. Long hair that extends to his back."

OUTPUT_DIR = "output/catboy04"

from torch import autocast
with autocast("cuda"):
    for i in range(30):
        image = pipe(prompt)["sample"][0]

        if not os.path.exists(OUTPUT_DIR):
            # ディレクトリが存在しない場合、ディレクトリを作成する
            os.makedirs(OUTPUT_DIR)

        image.save(f"{OUTPUT_DIR}/output{i:03}.png")