from diffusers import StableDiffusionPipeline
import torch
import random
import datetime

def random_seed(length):
    random.seed()
    min = 10**(length-1)
    max = 9*min + (min-1)
    return random.randint(min, max)

device_type = "cuda"  # Using AMD GPU with ROCm

def load_model():
    model_id = "./models/realistic-vision-v51"
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)

    pipe = pipe.to(device_type)
    return pipe

def generate_image(pipe, prompt, seed=None):
    generator = torch.Generator(device=device_type)

    if seed is not None:
        generator.manual_seed(seed)

    with torch.no_grad():
        image = pipe(prompt=prompt, num_inference_steps=20, guidance_scale=5, generator=generator).images[0]
    return image

quality_keywords = "Canon50, hyper detail, cinematic lighting, Canon EOS R3, nikon, f/1.4, ISO 200, 1/160s, 8K, RAW, unedited"

def generate(pipe, prompt, seed):
    image = generate_image(pipe, prompt, seed)
    dt = datetime.datetime.now().strftime("%y-%m-%d_%H-%M-%S")
    base_file = 'out/%s_%d' % (dt, seed)
    image_file = '%s.jpg' % base_file
    text_file = '%s.txt' % base_file

    image.save(image_file)
    with open(text_file, "w") as file:
        file.write("%d\n%s" % (seed, prompt))

    # Open in viewer
    image.show()


def main():
    pipe = load_model()
    prompt = "young adult woman, ((shoulder cut dark hair)), blue eyes, no makeup, white blouse, dressed, long sleeves, tiled head, enigmatic smile, %s" % quality_keywords

    seed = 673842166
    generate(pipe, prompt, seed)

if __name__ == "__main__":
    main()