import torch
import os
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel
from accelerate import PartialState
from diffusers import StableDiffusionPipeline
from diffusers import DiffusionPipeline

#model_path = "/home/gomishra/diffusers.old/examples/text_to_image/caleb_training_2"
#model_path ="/home/gomishra/Reliance/shareddata/reliance-model-lora-sdxl/"
model_path ="/shared/prerelease/home/gomishra/diffusers/examples/text_to_image/caleb_training"
#pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, variant="fp16",
#use_safetensors=True,)

pipe =DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)
#pipe.to("cuda")

refiner = DiffusionPipeline.from_pretrained(

    "stabilityai/stable-diffusion-xl-refiner-1.0",

    text_encoder_2=pipe.text_encoder_2,

    vae=pipe.vae,

    torch_dtype=torch.float16,

    use_safetensors=True,

    variant="fp16")

refiner.to("cuda")


prompts = {

    "amitabh bachchan":"amitabh bachchan in black suit with blue background and KBC as logo",

    "Prabhas":"prabhas with green background ",

    "Shah Rukh Khan":"Shah Rukh Khan on night market street",

    "Hritik Roshan":"Hritik Roshan singing on a stage at night "

}

folder_name = model_path.split("/")[-2]

#outDir = f"/data3/harshita_output/{folder_name}"

#outDir = f"/home/aac/sdxl_node2/output/try/{folder_name}"

outDir =f"/shared/prerelease/home/gomishra/diffusers/examples/text_to_image/outputdir"
if not os.path.exists(outDir):

    os.makedirs(outDir)

for key in list(prompts.keys()):

    print(key)

    prompt=prompts[key]

    image = pipe(

        prompt=prompt,

        num_inference_steps=50,

        denoising_end=0.8,

        guidance_scale=7.5,

        output_type="latent",

    ).images

    image = refiner(

        prompt=prompt,

        num_inference_steps=50,

        denoising_start=0.8,

        image=image,

    ).images[0]

    image.save(f"{outDir}/{key}.png")