import torch
from diffusers import StableDiffusionXLPipeline
import json
from tqdm import tqdm
import argparse

#pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16)
#pipe = pipe.to("cuda")

#prompt = "a photo of an astronaut riding a horse on mars"
#image = pipe(prompt).images[0]
#print(image)
def gen_image(part):
    pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, added_cond_kwargs={})
    pipe = pipe.to("cuda")


    data = json.load(open(data_path, "r"))
    data = data[part*800:(part+1)*800]
    for x in tqdm(data):
            prompt = "Generate an image for which the answer to this question: {} is {} based on the following image description: {}".format(x['src'], x['pred'], x['img_desc'])
            #prompt = "Generate an image for the following image description such that the answer to this question: Which breed of dog it this? is Affenpinscher: The image features a black dog in a grassy field, jumping in the air to catch a blue frisbee. The dog appears to be mid-leap, displaying its agility and athleticism. The frisbee is positioned in the air, close to the dog, indicating that it is the object of the dog's focus and effort."
            image = pipe(prompt).images[0]
            image.save("".format(x['id']))


if __name__ == "__main__":

    parser = argparse.ArgumentParser("Grounded-Segment-Anything Demo", add_help=True)
    parser.add_argument("--part", type=int, required=True, help="part")
    args = parser.parse_args()

    # cfg
    part = args.part  # change the path of the model config file
    gen_image(part)
