import json

from diffusers import DiffusionPipeline
import torch
import matplotlib.pyplot as plt


pipe = DiffusionPipeline.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0',
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant='fp16',
)

pipe.to('cuda')

with open('../examples/prompts.json', 'r') as f:
    prompts_list = list(json.load(f).values())

image_index = 4
prompt, negative_prompt = prompts_list[image_index]

images = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    negative_prompt=negative_prompt,
    seed=0,
    num_inference_steps=20,
    guidance_scale=9.0,
).images[0]
images.save(fr'C:\Users\dragynir\Pictures\FASHION_IMAGES\no_condition\image_{image_index+1}.png')
plt.imshow(images)
plt.show()