import numpy as np
from PIL import Image
import cv2
import torch

count = 0
import torch
import numpy as np
from PIL import Image, ExifTags
from einops import rearrange
from transformers import pipeline
import tqdm
import random
from pathlib import Path


from flux.cli import SamplingOptions
from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
from flux.util import configs, embed_watermark, load_ae, load_clip, load_flow_model, load_t5

def get_models(name: str, device: torch.device, offload: bool, is_schnell: bool):
    t5 = load_t5(device, max_length=256 if is_schnell else 512)
    clip = load_clip(device)
    model = load_flow_model(name, device="cpu" if offload else device)
    ae = load_ae(name, device="cpu" if offload else device)
    return model, ae, t5, clip

class FluxImageToImage:
    def __init__(self, model_name: str, devices: list = None, offload: bool = False):
        self.devices = devices if devices else ["cuda:0", "cuda:1"]
        self.offload = offload
        self.model_name = model_name
        self.is_schnell = model_name == "flux-schnell"
        self.models = [get_models(model_name, torch.device(d), offload, self.is_schnell) for d in self.devices]

    @torch.inference_mode()
    def generate(
        self,
        init_images,
        prompts,
        width=1024,
        height=720,
        num_steps=40,
        guidance=3.5,
        seeds=None,
        image2image_strength=0.9,
        add_sampling_metadata=True,
    ):
        batch_size = len(init_images)
        seeds = seeds or [-1] * batch_size
        opts = [SamplingOptions(prompt=prompts[i], width=width, height=height, num_steps=num_steps, guidance=guidance, seed=seeds[i]) for i in range(batch_size)]

        # Fix for random seed generation
        seeds = [torch.Generator(device="cpu").seed() if seed == -1 else seed for seed in seeds]

        results = []
        device_batches = [[] for _ in self.devices]

        # Distribute init_images and options across devices
        for idx in range(batch_size):
            device_idx = idx % len(self.devices)
            device_batches[device_idx].append((init_images[idx], opts[idx]))

        for i, device in enumerate(self.devices):
            model, ae, t5, clip = self.models[i]
            device = torch.device(device)
            
            for init_image, opt in device_batches[i]:
                init_image = init_image.to(device)
                init_image = torch.nn.functional.interpolate(init_image, (opt.height, opt.width))
                
                if self.offload:
                    ae.encoder.to(device)
                init_image = ae.encode(init_image.to(device))
                if self.offload:
                    ae = ae.cpu()
                    torch.cuda.empty_cache()

                x = get_noise(
                    1,
                    opt.height,
                    opt.width,
                    device=device,
                    dtype=torch.bfloat16,
                    seed=opt.seed,
                )
                timesteps = get_schedule(
                    opt.num_steps,
                    x.shape[-1] * x.shape[-2] // 4,
                    shift=(not self.is_schnell),
                )
                t_idx = int((1 - image2image_strength) * num_steps)
                t = timesteps[t_idx]
                timesteps = timesteps[t_idx:]
                x = t * x + (1.0 - t) * init_image.to(x.dtype)

                if self.offload:
                    t5, clip = t5.to(device), clip.to(device)
                inp = prepare(t5=t5, clip=clip, img=x, prompt=opt.prompt)

                if self.offload:
                    t5, clip = t5.cpu(), clip.cpu()
                    torch.cuda.empty_cache()
                    model = model.to(device)

                x = denoise(model, **inp, timesteps=timesteps, guidance=opt.guidance)

                if self.offload:
                    model.cpu()
                    torch.cuda.empty_cache()
                    ae.decoder.to(x.device)

                x = unpack(x.float(), opt.height, opt.width)
                with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
                    x = ae.decode(x)

                if self.offload:
                    ae.decoder.cpu()
                    torch.cuda.empty_cache()

                x = x.clamp(-1, 1)
                x = embed_watermark(x.float())
                x = rearrange(x[0], "c h w -> h w c")

                img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
                # nsfw_score = 0.7
                # if nsfw_score < 0.85:
                #     exif_data = Image.Exif()
                #     exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
                #     exif_data[ExifTags.Base.Make] = "Black Forest Labs"
                #     exif_data[ExifTags.Base.Model] = self.model_name
                #     if add_sampling_metadata:
                #         exif_data[ExifTags.Base.ImageDescription] = opt.prompt
                    
                output_file = f"output_{i}.jpg"
                #     # img.save(output_file, format="jpeg", exif=exif_data, quality=95, subsampling=0)
                results.append((img, str(opt.seed), output_file, None))
                # else:
                #     results.append((None, str(opt.seed), None, "Your generated image may contain NSFW content."))

        return results

def generate_prompt():
    ethnicities = ["Asian", "African", "European", "Middle Eastern", "Latin American", "South Asian", "Southeast Asian", "Pacific Islander"]
    skin_tones = ["fair", "light", "medium", "olive", "brown", "dark"]
    body_types = ["slim", "athletic", "curvy", "plus-size", "muscular", "petite"]
    
    gender = random.choice(["male", "female"])
    ethnicity = random.choice(ethnicities)
    skin_tone = random.choice(skin_tones)
    body_type = random.choice(body_types)
    clothing = random.choice(["traditional", "modern", "casual", "formal", "ethnic", "professional", "sporty"])
    colors = ["red", "blue", "green", "yellow", "purple", "orange", "white", "black", "pink", "teal", "maroon", "navy"]
    color = random.choice(colors)

    if gender == "female":
        outfit = random.choice(["sari", "lehenga", "salwar kameez", "dress", "skirt and top", "jeans and blouse", "pantsuit"])
        hair = random.choice(["long", "short", "curly", "straight", "wavy", "braided", "updone"])
        accessory = random.choice(["necklace", "earrings", "bangles", "anklets", "bindi", "scarf", "handbag"])
    else:
        outfit = random.choice(["kurta pajama", "sherwani", "dhoti", "shirt and pants", "suit", "jeans and t-shirt", "traditional ethnic wear"])
        hair = random.choice(["short", "medium-length", "styled", "messy", "neat", "bald", "beard"])
        accessory = random.choice(["turban", "pagri", "watch", "scarf", "sunglasses", "tie", "belt"])

    prompt = f"A character sheet of a {gender} {ethnicity} person with {skin_tone} skin and a {body_type} build, "
    prompt += f"wearing a {color} {clothing} {outfit} with {hair} hair, "
    prompt += f"adorned with a {accessory}. In different poses and angles, including front view, side view, back view, and three-quarter view."

    return prompt

def generate_dataset(total_images, batch_size, flux_img2img):
    global count
    output_dir = Path("/ihub/homedirs/vp_hc/neuralnexus/InfiniGPT/dataset/output_batched")
    output_dir.mkdir(exist_ok=True)

    for _ in tqdm.tqdm(range(total_images // batch_size)):
        prompts = [generate_prompt() for _ in range(batch_size)]
        init_images = [Image.open("/ihub/homedirs/vp_hc/neuralnexus/InfiniGPT/dataset/reference2.jpg") for _ in range(batch_size)]
        init_images = [torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0 for img in init_images]
        init_images = [img.unsqueeze(0) for img in init_images]

        results = flux_img2img.generate(
            init_images=init_images,
            prompts=prompts,
            width=1024,
            height=720,
            num_steps=40,
            guidance=3.5,
            seeds=[-1] * batch_size,
            image2image_strength=0.9
        )

        for i, (result_image, seed, output_path, warning) in enumerate(results):
            if result_image:
                output_file = output_dir / f"output_{count:06d}.jpg"
                result_image.save(output_file)
                count +=1
        print(f"Total images generated: {count}")

if __name__ == "__main__":
    flux_img2img = FluxImageToImage(model_name="flux-dev", devices=["cuda"])
    generate_dataset(total_images=100000, batch_size=25, flux_img2img=flux_img2img)

the above code is wrong for batched generation. Please correct it properly its not working as intended