tools/synth.py

"""
Helper scripts for generating synthetic images using diffusion model.

Functions:
    - get_top_misclassified
    - get_class_list
    - generateClassPairs
    - outputDirectory
    - pipe_img
    - createPrompts
    - interpolatePrompts
        - slerp
        - get_middle_elements
        - remove_middle
    - genClassImg
    - getMetadata
    - groupbyInterpolation
    - ungroupInterpolation
    - groupAllbyInterpolation
    - getPairIndices
    - generateImagesFromDataset
    - generateTrace
"""

import json
import os

import numpy as np
import pandas as pd
import torch
from DeepCache import DeepCacheSDHelper
from diffusers import (
    LMSDiscreteScheduler,
    StableDiffusionImg2ImgPipeline,
)
from torch import nn
from torchmetrics.functional.image import structural_similarity_index_measure as ssim
from torchvision import transforms


def get_top_misclassified(val_classifier_json):
    """
    Retrieves the top misclassified classes from a validation classifier JSON file.

    Args:
        val_classifier_json (str): The path to the validation classifier JSON file.

    Returns:
        dict: A dictionary containing the top misclassified classes, where the keys are the class names
              and the values are the number of misclassifications.
    """
    with open(val_classifier_json) as f:
        val_output = json.load(f)
    val_metrics_df = pd.DataFrame.from_dict(
        val_output["val_metrics_details"], orient="index"
    )
    class_dict = dict()
    for k, v in val_metrics_df["top_n_classes"].items():
        class_dict[k] = v
    return class_dict


def get_class_list(val_classifier_json):
    """
    Retrieves the list of classes from the given validation classifier JSON file.

    Args:
        val_classifier_json (str): The path to the validation classifier JSON file.

    Returns:
        list: A sorted list of class names extracted from the JSON file.
    """
    with open(val_classifier_json, "r") as f:
        data = json.load(f)
    return sorted(list(data["val_metrics_details"].keys()))


def generateClassPairs(val_classifier_json):
    """
    Generate pairs of misclassified classes from the given validation classifier JSON.

    Args:
        val_classifier_json (str): The path to the validation classifier JSON file.

    Returns:
        list: A sorted list of pairs of misclassified classes.
    """
    pairs = set()
    misclassified_classes = get_top_misclassified(val_classifier_json)
    for key, value in misclassified_classes.items():
        for v in value:
            pairs.add(tuple(sorted([key, v])))
    return sorted(list(pairs))


def outputDirectory(class_pairs, synth_path, metadata_path):
    """
    Creates the output directory structure for the synthesized data.

    Args:
        class_pairs (list): A list of class pairs.
        synth_path (str): The path to the directory where the synthesized data will be stored.
        metadata_path (str): The path to the directory where the metadata will be stored.

    Returns:
        None
    """
    for id in class_pairs:
        class_folder = f"{synth_path}/{id}"
        if not (os.path.exists(class_folder)):
            os.makedirs(class_folder)
    if not (os.path.exists(metadata_path)):
        os.makedirs(metadata_path)
    print("Info: Output directory ready.")


def pipe_img(
    model_path,
    device="cuda",
    apply_optimization=True,
    use_torchcompile=False,
    ci_cb=(5, 1),
    use_safetensors=None,
    cpu_offload=False,
    scheduler=None,
):
    """
    Creates and returns an image-to-image pipeline for stable diffusion.

    Args:
        model_path (str): The path to the pretrained model.
        device (str, optional): The device to use for computation. Defaults to "cuda".
        apply_optimization (bool, optional): Whether to apply optimization techniques. Defaults to True.
        use_torchcompile (bool, optional): Whether to use torchcompile for model compilation. Defaults to False.
        ci_cb (tuple, optional): A tuple containing the cache interval and cache branch ID. Defaults to (5, 1).
        use_safetensors (bool, optional): Whether to use safetensors. Defaults to None.
        cpu_offload (bool, optional): Whether to enable CPU offloading. Defaults to False.
        scheduler (LMSDiscreteScheduler, optional): The scheduler for the pipeline. Defaults to None.

    Returns:
        StableDiffusionImg2ImgPipeline: The image-to-image pipeline for stable diffusion.
    """
    ###############################
    # Reference:
    # Akimov, R. (2024) Images Interpolation with Stable Diffusion - Hugging Face Open-Source AI Cookbook. Available at: https://huggingface.co/learn/cookbook/en/stable_diffusion_interpolation (Accessed: 4 June 2024).
    ###############################
    if scheduler is None:
        scheduler = LMSDiscreteScheduler(
            beta_start=0.00085,
            beta_end=0.012,
            beta_schedule="scaled_linear",
            num_train_timesteps=1000,
            steps_offset=1,
        )
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model_path,
        scheduler=scheduler,
        torch_dtype=torch.float32,
        use_safetensors=use_safetensors,
    ).to(device)
    if cpu_offload:
        pipe.enable_model_cpu_offload()
    if apply_optimization:
        # tomesd.apply_patch(pipe, ratio=0.5)
        helper = DeepCacheSDHelper(pipe=pipe)
        cache_interval, cache_branch_id = ci_cb
        helper.set_params(
            cache_interval=cache_interval, cache_branch_id=cache_branch_id
        )  # lower is faster but lower quality
        helper.enable()
        if torch.cuda.is_available():
            pipe.enable_xformers_memory_efficient_attention()
        if use_torchcompile:
            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
    return pipe


def createPrompts(
    class_name_pairs,
    prompt_structure=None,
    use_default_negative_prompt=False,
    negative_prompt=None,
):
    """
    Create prompts for image generation.

    Args:
        class_name_pairs (list): A list of two class names.
        prompt_structure (str, optional): The structure of the prompt. Defaults to "a photo of a <class_name>".
        use_default_negative_prompt (bool, optional): Whether to use the default negative prompt. Defaults to False.
        negative_prompt (str, optional): The negative prompt to steer the generation away from certain features.

    Returns:
        tuple: A tuple containing two lists - prompts and negative_prompts.
            prompts (list): Text prompts that describe the desired output image.
            negative_prompts (list): Negative prompts that can be used to steer the generation away from certain features.
    """
    if prompt_structure is None:
        prompt_structure = "a photo of a <class_name>"
    elif "<class_name>" not in prompt_structure:
        raise ValueError(
            "The prompt structure must contain the <class_name> placeholder."
        )
    if use_default_negative_prompt:
        default_negative_prompt = (
            "blurry image, disfigured, deformed, distorted, cartoon, drawings"
        )
        negative_prompt = default_negative_prompt

    class1 = class_name_pairs[0]
    class2 = class_name_pairs[1]
    prompt1 = prompt_structure.replace("<class_name>", class1)
    prompt2 = prompt_structure.replace("<class_name>", class2)
    prompts = [prompt1, prompt2]
    if negative_prompt is None:
        print("Info: Negative prompt not provided, returning as None.")
        return prompts, None
    else:
        # Negative prompts that can be used to steer the generation away from certain features.
        negative_prompts = [negative_prompt] * len(prompts)
        return prompts, negative_prompts


def interpolatePrompts(
    prompts,
    pipeline,
    num_interpolation_steps,
    sample_mid_interpolation,
    remove_n_middle=0,
    device="cuda",
):
    """
    Interpolates prompts by generating intermediate embeddings between pairs of prompts.

    Args:
        prompts (List[str]): A list of prompts to be interpolated.
        pipeline: The pipeline object containing the tokenizer and text encoder.
        num_interpolation_steps (int): The number of interpolation steps between each pair of prompts.
        sample_mid_interpolation (int): The number of intermediate embeddings to sample from the middle of the interpolated prompts.
        remove_n_middle (int, optional): The number of middle embeddings to remove from the interpolated prompts. Defaults to 0.
        device (str, optional): The device to run the interpolation on. Defaults to "cuda".

    Returns:
        interpolated_prompt_embeds (torch.Tensor): The interpolated prompt embeddings.
        prompt_metadata (dict): Metadata about the interpolation process, including similarity scores and nearest class information.

    e.g. if num_interpolation_steps = 10, sample_mid_interpolation = 6, remove_n_middle = 2
    Interpolated: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    Sampled:            [2, 3, 4, 5, 6, 7]
    Removed:                   x  x
    Returns:            [2, 3,       6, 7]
    """

    ###############################
    # Reference:
    # Akimov, R. (2024) Images Interpolation with Stable Diffusion - Hugging Face Open-Source AI Cookbook. Available at: https://huggingface.co/learn/cookbook/en/stable_diffusion_interpolation (Accessed: 4 June 2024).
    ###############################

    def slerp(v0, v1, num, t0=0, t1=1):
        """
        Performs spherical linear interpolation between two vectors.

        Args:
            v0 (torch.Tensor): The starting vector.
            v1 (torch.Tensor): The ending vector.
            num (int): The number of interpolation points.
            t0 (float, optional): The starting time. Defaults to 0.
            t1 (float, optional): The ending time. Defaults to 1.

        Returns:
            torch.Tensor: The interpolated vectors.

        """
        ###############################
        # Reference:
        # Karpathy, A. (2022) hacky stablediffusion code for generating videos, Gist. Available at: https://gist.github.com/karpathy/00103b0037c5aaea32fe1da1af553355 (Accessed: 4 June 2024).
        ###############################
        v0 = v0.detach().cpu().numpy()
        v1 = v1.detach().cpu().numpy()

        def interpolation(t, v0, v1, DOT_THRESHOLD=0.9995):
            """helper function to spherically interpolate two arrays v1 v2"""
            dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
            if np.abs(dot) > DOT_THRESHOLD:
                v2 = (1 - t) * v0 + t * v1
            else:
                theta_0 = np.arccos(dot)
                sin_theta_0 = np.sin(theta_0)
                theta_t = theta_0 * t
                sin_theta_t = np.sin(theta_t)
                s0 = np.sin(theta_0 - theta_t) / sin_theta_0
                s1 = sin_theta_t / sin_theta_0
                v2 = s0 * v0 + s1 * v1
            return v2

        t = np.linspace(t0, t1, num)

        v3 = torch.tensor(np.array([interpolation(t[i], v0, v1) for i in range(num)]))

        return v3

    def get_middle_elements(lst, n):
        """
        Returns a tuple containing a sublist of the middle elements of the given list `lst` and a range of indices of those elements.

        Args:
            lst (list): The list from which to extract the middle elements.
            n (int): The number of middle elements to extract.

        Returns:
            tuple: A tuple containing the sublist of middle elements and a range of indices.

        Raises:
            None

        Examples:
            lst = [1, 2, 3, 4, 5]
            get_middle_elements(lst, 3)
            ([2, 3, 4], range(2, 5))
        """
        if n % 2 == 0:  # Even number of elements
            middle_index = len(lst) // 2 - 1
            start = middle_index - n // 2 + 1
            end = middle_index + n // 2 + 1
            return lst[start:end], range(start, end)
        else:  # Odd number of elements
            middle_index = len(lst) // 2
            start = middle_index - n // 2
            end = middle_index + n // 2 + 1
            return lst[start:end], range(start, end)

    def remove_middle(data, n):
        """
        Remove the middle n elements from a list.

        Args:
            data (list): The input list.
            n (int): The number of elements to remove from the middle of the list.

        Returns:
            list: The modified list with the middle n elements removed.

        Raises:
            ValueError: If n is negative or greater than the length of the list.

        """
        if n < 0 or n > len(data):
            raise ValueError(
                "Invalid value for n. It should be non-negative and less than half the list length"
            )

        # Find the middle index
        middle = len(data) // 2

        # Create slices to exclude the middle n elements
        if n == 1:
            return data[:middle] + data[middle + 1 :]
        elif n % 2 == 0:
            return data[: middle - n // 2] + data[middle + n // 2 :]
        else:
            return data[: middle - n // 2] + data[middle + n // 2 + 1 :]

    batch_size = len(prompts)

    # Tokenizing and encoding prompts into embeddings.
    prompts_tokens = pipeline.tokenizer(
        prompts,
        padding="max_length",
        max_length=pipeline.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    prompts_embeds = pipeline.text_encoder(prompts_tokens.input_ids.to(device))[0]

    # Interpolating between embeddings pairs for the given number of interpolation steps.
    interpolated_prompt_embeds = []

    for i in range(batch_size - 1):
        interpolated_prompt_embeds.append(
            slerp(prompts_embeds[i], prompts_embeds[i + 1], num_interpolation_steps)
        )

    full_interpolated_prompt_embeds = interpolated_prompt_embeds[:]
    interpolated_prompt_embeds[0], sample_range = get_middle_elements(
        interpolated_prompt_embeds[0], sample_mid_interpolation
    )

    if remove_n_middle > 0:
        interpolated_prompt_embeds[0] = remove_middle(
            interpolated_prompt_embeds[0], remove_n_middle
        )

    prompt_metadata = dict()
    similarity = nn.CosineSimilarity(dim=-1, eps=1e-6)
    for i in range(num_interpolation_steps):
        class1_sim = (
            similarity(
                full_interpolated_prompt_embeds[0][0],
                full_interpolated_prompt_embeds[0][i],
            )
            .mean()
            .item()
        )
        class2_sim = (
            similarity(
                full_interpolated_prompt_embeds[0][num_interpolation_steps - 1],
                full_interpolated_prompt_embeds[0][i],
            )
            .mean()
            .item()
        )
        relative_distance = class1_sim / (class1_sim + class2_sim)

        prompt_metadata[i] = {
            "selected": i in sample_range,
            "similarity": {
                "class1": class1_sim,
                "class2": class2_sim,
                "class1_relative_distance": relative_distance,
                "class2_relative_distance": 1 - relative_distance,
            },
            "nearest_class": int(relative_distance < 0.5),
        }

    interpolated_prompt_embeds = torch.cat(interpolated_prompt_embeds, dim=0).to(device)
    return interpolated_prompt_embeds, prompt_metadata


def genClassImg(
    pipeline,
    pos_embed,
    neg_embed,
    input_image,
    generator,
    latents,
    num_imgs=1,
    height=512,
    width=512,
    num_inference_steps=25,
    guidance_scale=7.5,
):
    """
    Generate class image using the given inputs.

    Args:
        pipeline: The pipeline object used for image generation.
        pos_embed: The positive embedding for the class.
        neg_embed: The negative embedding for the class (optional).
        input_image: The input image for guidance (optional).
        generator: The generator model used for image generation.
        latents: The latent vectors used for image generation.
        num_imgs: The number of images to generate (default is 1).
        height: The height of the generated images (default is 512).
        width: The width of the generated images (default is 512).
        num_inference_steps: The number of inference steps for image generation (default is 25).
        guidance_scale: The scale factor for guidance (default is 7.5).

    Returns:
        The generated class image.
    """

    if neg_embed is not None:
        npe = neg_embed[None, ...]
    else:
        npe = None

    return pipeline(
        height=height,
        width=width,
        num_images_per_prompt=num_imgs,
        prompt_embeds=pos_embed[None, ...],
        negative_prompt_embeds=npe,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
        latents=latents,
        image=input_image,
    ).images[0]


def getMetadata(
    class_pairs,
    path,
    seed,
    guidance_scale,
    num_inference_steps,
    num_interpolation_steps,
    sample_mid_interpolation,
    height,
    width,
    prompts,
    negative_prompts,
    pipeline,
    prompt_metadata,
    negative_prompt_metadata,
    ssim_metadata=None,
    save_json=True,
    save_path=".",
):
    """
    Generate metadata for the given parameters.

    Args:
        class_pairs (list): List of class pairs.
        path (str): Path to the data.
        seed (int): Seed value for randomization.
        guidance_scale (float): Scale factor for guidance.
        num_inference_steps (int): Number of inference steps.
        num_interpolation_steps (int): Number of interpolation steps.
        sample_mid_interpolation (bool): Flag to sample mid-interpolation.
        height (int): Height of the image.
        width (int): Width of the image.
        prompts (list): List of prompts.
        negative_prompts (list): List of negative prompts.
        pipeline (object): Pipeline object.
        prompt_metadata (dict): Metadata for prompts.
        negative_prompt_metadata (dict): Metadata for negative prompts.
        ssim_metadata (dict, optional): SSIM scores metadata. Defaults to None.
        save_json (bool, optional): Flag to save metadata as JSON. Defaults to True.
        save_path (str, optional): Path to save the JSON file. Defaults to ".".

    Returns:
        dict: Generated metadata.
    """

    metadata = dict()

    metadata["class_pairs"] = class_pairs
    metadata["path"] = path
    metadata["seed"] = seed
    metadata["params"] = {
        "CFG": guidance_scale,
        "inferenceSteps": num_inference_steps,
        "interpolationSteps": num_interpolation_steps,
        "sampleMidInterpolation": sample_mid_interpolation,
        "height": height,
        "width": width,
    }
    for i in range(len(prompts)):
        metadata[f"prompt_text_{i}"] = prompts[i]
        if negative_prompts is not None:
            metadata[f"negative_prompt_text_{i}"] = negative_prompts[i]
    metadata["pipe_config"] = dict(pipeline.config)
    metadata["prompt_embed_similarity"] = prompt_metadata
    metadata["negative_prompt_embed_similarity"] = negative_prompt_metadata
    if ssim_metadata is not None:
        print("Info: SSIM scores are available.")
        metadata["ssim_scores"] = ssim_metadata
    if save_json:
        with open(
            os.path.join(save_path, f"{'_'.join(i for i in class_pairs)}_{seed}.json"),
            "w",
        ) as f:
            json.dump(metadata, f, indent=4)
    return metadata


def groupbyInterpolation(dir_to_classfolder):
    """
    Group files in a directory by interpolation step.

    Args:
        dir_to_classfolder (str): The path to the directory containing the files.

    Returns:
        None
    """
    files = [
        (f.split(sep="_")[1].split(sep=".")[0], os.path.join(dir_to_classfolder, f))
        for f in os.listdir(dir_to_classfolder)
    ]
    # create a subfolder for each step of the interpolation
    for interpolation_step, file_path in files:
        new_dir = os.path.join(dir_to_classfolder, interpolation_step)
        if not os.path.exists(new_dir):
            os.makedirs(new_dir)
        os.rename(file_path, os.path.join(new_dir, os.path.basename(file_path)))


def ungroupInterpolation(dir_to_classfolder):
    """
    Moves all files from subdirectories within `dir_to_classfolder` to `dir_to_classfolder` itself,
    and then removes the subdirectories.

    Args:
        dir_to_classfolder (str): The path to the directory containing the subdirectories.

    Returns:
        None
    """
    for interpolation_step in os.listdir(dir_to_classfolder):
        if os.path.isdir(os.path.join(dir_to_classfolder, interpolation_step)):
            for f in os.listdir(os.path.join(dir_to_classfolder, interpolation_step)):
                os.rename(
                    os.path.join(dir_to_classfolder, interpolation_step, f),
                    os.path.join(dir_to_classfolder, f),
                )
            os.rmdir(os.path.join(dir_to_classfolder, interpolation_step))


def groupAllbyInterpolation(
    data_path,
    group=True,
    fn_group=groupbyInterpolation,
    fn_ungroup=ungroupInterpolation,
):
    """
    Group or ungroup all data classes by interpolation.

    Args:
        data_path (str): The path to the data.
        group (bool, optional): Whether to group the data. Defaults to True.
        fn_group (function, optional): The function to use for grouping. Defaults to groupbyInterpolation.
        fn_ungroup (function, optional): The function to use for ungrouping. Defaults to ungroupInterpolation.
    """
    data_classes = sorted(os.listdir(data_path))
    if group:
        fn = fn_group
    else:
        fn = fn_ungroup
    for c in data_classes:
        c_path = os.path.join(data_path, c)
        if os.path.isdir(c_path):
            fn(c_path)
            print(f"Processed {c}")


def getPairIndices(subset_len, total_pair_count=1, seed=None):
    """
    Generate pairs of indices for a given subset length.

    Args:
        subset_len (int): The length of the subset.
        total_pair_count (int, optional): The total number of pairs to generate. Defaults to 1.
        seed (int, optional): The seed value for the random number generator. Defaults to None.

    Returns:
        list: A list of pairs of indices.

    """
    rng = np.random.default_rng(seed)
    group_size = (subset_len + total_pair_count - 1) // total_pair_count
    numbers = list(range(subset_len))
    numbers_selection = list(range(subset_len))
    rng.shuffle(numbers)
    for i in range(group_size - subset_len % group_size):
        numbers.append(numbers_selection[i])
    numbers = np.array(numbers)
    groups = numbers[: group_size * total_pair_count].reshape(-1, group_size)
    return groups.tolist()


def generateImagesFromDataset(
    img_subsets,
    class_iterables,
    pipeline,
    interpolated_prompt_embeds,
    interpolated_negative_prompts_embeds,
    num_inference_steps,
    guidance_scale,
    height=512,
    width=512,
    seed=None,
    save_path=".",
    class_pairs=("0", "1"),
    save_image=True,
    image_type="jpg",
    interpolate_range="full",
    device="cuda",
    return_images=False,
):
    """
    Generates images from a dataset using the given parameters.

    Args:
        img_subsets (dict): A dictionary containing image subsets for each class.
        class_iterables (dict): A dictionary containing iterable objects for each class.
        pipeline (object): The pipeline object used for image generation.
        interpolated_prompt_embeds (list): A list of interpolated prompt embeddings.
        interpolated_negative_prompts_embeds (list): A list of interpolated negative prompt embeddings.
        num_inference_steps (int): The number of inference steps for image generation.
        guidance_scale (float): The scale factor for guidance loss during image generation.
        height (int, optional): The height of the generated images. Defaults to 512.
        width (int, optional): The width of the generated images. Defaults to 512.
        seed (int, optional): The seed value for random number generation. Defaults to None.
        save_path (str, optional): The path to save the generated images. Defaults to ".".
        class_pairs (tuple, optional): A tuple containing pairs of class identifiers. Defaults to ("0", "1").
        save_image (bool, optional): Whether to save the generated images. Defaults to True.
        image_type (str, optional): The file format of the saved images. Defaults to "jpg".
        interpolate_range (str, optional): The range of interpolation for prompt embeddings.
            Possible values are "full", "nearest", or "furthest". Defaults to "full".
        device (str, optional): The device to use for image generation. Defaults to "cuda".
        return_images (bool, optional): Whether to return the generated images. Defaults to False.

    Returns:
        dict or tuple: If return_images is True, returns a dictionary containing the generated images for each class and a dictionary containing the SSIM scores for each class and interpolation step.
                       If return_images is False, returns a dictionary containing the SSIM scores for each class and interpolation step.
    """
    if interpolate_range == "nearest":
        nearest_half = True
        furthest_half = False
    elif interpolate_range == "furthest":
        nearest_half = False
        furthest_half = True
    else:
        nearest_half = False
        furthest_half = False

    if seed is None:
        seed = torch.Generator().seed()
    generator = torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    # Generating initial U-Net latent vectors from a random normal distribution.
    latents = torch.randn(
        (1, pipeline.unet.config.in_channels, height // 8, width // 8),
        generator=generator,
    ).to(device)

    embed_len = len(interpolated_prompt_embeds)
    embed_pairs = zip(interpolated_prompt_embeds, interpolated_negative_prompts_embeds)
    embed_pairs_list = list(embed_pairs)
    if return_images:
        class_images = dict()
    class_ssim = dict()

    if nearest_half or furthest_half:
        if nearest_half:
            steps_range = (range(0, embed_len // 2), range(embed_len // 2, embed_len))
            mutiplier = 2
        elif furthest_half:
            # uses opposite class of images of the text interpolation
            steps_range = (range(embed_len // 2, embed_len), range(0, embed_len // 2))
            mutiplier = 2
    else:
        steps_range = (range(embed_len), range(embed_len))
        mutiplier = 1

    for class_iter, class_id in enumerate(class_pairs):
        if return_images:
            class_images[class_id] = list()
        class_ssim[class_id] = {
            i: {"ssim_sum": 0, "ssim_count": 0, "ssim_avg": 0} for i in range(embed_len)
        }
        subset_len = len(img_subsets[class_id])
        # to efficiently randomize the steps to interpolate for each image in the class, group_map is used
        # group_map: index is the image id, element is the group id
        # steps_range[class_iter] determines the range of steps to interpolate for the class,
        # so the first half of the steps are for the first class and so on. range(0,7) and range(8,15) for 16 steps
        # then the rest is to multiply the steps to cover the whole subset + remainder
        group_map = (
            list(steps_range[class_iter]) * mutiplier * (subset_len // embed_len + 1)
        )
        rng.shuffle(
            group_map
        )  # shuffle the steps to interpolate for each image, position in the group_map is mapped to the image id

        iter_indices = class_iterables[class_id].pop()
        # generate images for each image in the class, randomly selecting an interpolated step
        for image_id in iter_indices:
            img, trg = img_subsets[class_id][image_id]
            input_image = img.unsqueeze(0)
            interpolate_step = group_map[image_id]
            prompt_embeds, negative_prompt_embeds = embed_pairs_list[interpolate_step]
            generated_image = genClassImg(
                pipeline,
                prompt_embeds,
                negative_prompt_embeds,
                input_image,
                generator,
                latents,
                num_imgs=1,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
            )
            pred_image = transforms.ToTensor()(generated_image).unsqueeze(0)
            ssim_score = ssim(pred_image, input_image).item()
            class_ssim[class_id][interpolate_step]["ssim_sum"] += ssim_score
            class_ssim[class_id][interpolate_step]["ssim_count"] += 1
            if return_images:
                class_images[class_id].append(generated_image)
            if save_image:
                if image_type == "jpg":
                    generated_image.save(
                        f"{save_path}/{class_id}/{seed}-{image_id}_{interpolate_step}.{image_type}",
                        format="JPEG",
                        quality=95,
                    )
                elif image_type == "png":
                    generated_image.save(
                        f"{save_path}/{class_id}/{seed}-{image_id}_{interpolate_step}.{image_type}",
                        format="PNG",
                    )
                else:
                    generated_image.save(
                        f"{save_path}/{class_id}/{seed}-{image_id}_{interpolate_step}.{image_type}"
                    )

        # calculate ssim avg for the class
        for i_step in range(embed_len):
            if class_ssim[class_id][i_step]["ssim_count"] > 0:
                class_ssim[class_id][i_step]["ssim_avg"] = (
                    class_ssim[class_id][i_step]["ssim_sum"]
                    / class_ssim[class_id][i_step]["ssim_count"]
                )

    if return_images:
        return class_images, class_ssim
    else:
        return class_ssim


def generateTrace(
    prompts,
    img_subsets,
    class_iterables,
    interpolated_prompt_embeds,
    interpolated_negative_prompts_embeds,
    subset_indices,
    seed=None,
    save_path=".",
    class_pairs=("0", "1"),
    image_type="jpg",
    interpolate_range="full",
    save_prompt_embeds=False,
):
    """
    Generate a trace dictionary containing information about the generated images.

    Args:
        prompts (list): List of prompt texts.
        img_subsets (dict): Dictionary containing image subsets for each class.
        class_iterables (dict): Dictionary containing iterable objects for each class.
        interpolated_prompt_embeds (torch.Tensor): Tensor containing interpolated prompt embeddings.
        interpolated_negative_prompts_embeds (torch.Tensor): Tensor containing interpolated negative prompt embeddings.
        subset_indices (dict): Dictionary containing indices of subsets for each class.
        seed (int, optional): Seed value for random number generation. Defaults to None.
        save_path (str, optional): Path to save the generated images. Defaults to ".".
        class_pairs (tuple, optional): Tuple containing class pairs. Defaults to ("0", "1").
        image_type (str, optional): Type of the generated images. Defaults to "jpg".
        interpolate_range (str, optional): Range of interpolation. Defaults to "full".
        save_prompt_embeds (bool, optional): Flag to save prompt embeddings. Defaults to False.

    Returns:
        dict: Trace dictionary containing information about the generated images.
    """
    trace_dict = {
        "class_pairs": list(),
        "class_id": list(),
        "image_id": list(),
        "interpolation_step": list(),
        "embed_len": list(),
        "pos_prompt_text": list(),
        "neg_prompt_text": list(),
        "input_file_path": list(),
        "output_file_path": list(),
        "input_prompts_embed": list(),
    }

    if interpolate_range == "nearest":
        nearest_half = True
        furthest_half = False
    elif interpolate_range == "furthest":
        nearest_half = False
        furthest_half = True
    else:
        nearest_half = False
        furthest_half = False

    if seed is None:
        seed = torch.Generator().seed()
    rng = np.random.default_rng(seed)

    embed_len = len(interpolated_prompt_embeds)
    embed_pairs = zip(
        interpolated_prompt_embeds.cpu().numpy(),
        interpolated_negative_prompts_embeds.cpu().numpy(),
    )
    embed_pairs_list = list(embed_pairs)

    if nearest_half or furthest_half:
        if nearest_half:
            steps_range = (range(0, embed_len // 2), range(embed_len // 2, embed_len))
            mutiplier = 2
        elif furthest_half:
            # uses opposite class of images of the text interpolation
            steps_range = (range(embed_len // 2, embed_len), range(0, embed_len // 2))
            mutiplier = 2
    else:
        steps_range = (range(embed_len), range(embed_len))
        mutiplier = 1

    for class_iter, class_id in enumerate(class_pairs):

        subset_len = len(img_subsets[class_id])
        # to efficiently randomize the steps to interpolate for each image in the class, group_map is used
        # group_map: index is the image id, element is the group id
        # steps_range[class_iter] determines the range of steps to interpolate for the class,
        # so the first half of the steps are for the first class and so on. range(0,7) and range(8,15) for 16 steps
        # then the rest is to multiply the steps to cover the whole subset + remainder
        group_map = (
            list(steps_range[class_iter]) * mutiplier * (subset_len // embed_len + 1)
        )
        rng.shuffle(
            group_map
        )  # shuffle the steps to interpolate for each image, position in the group_map is mapped to the image id

        iter_indices = class_iterables[class_id].pop()
        # generate images for each image in the class, randomly selecting an interpolated step
        for image_id in iter_indices:
            class_ds = img_subsets[class_id]
            interpolate_step = group_map[image_id]
            sample_count = subset_indices[class_id][0] + image_id
            input_file = os.path.normpath(class_ds.dataset.samples[sample_count][0])
            pos_prompt = prompts[0]
            neg_prompt = prompts[1]
            output_file = f"{save_path}/{class_id}/{seed}-{image_id}_{interpolate_step}.{image_type}"
            if save_prompt_embeds:
                input_prompts_embed = embed_pairs_list[interpolate_step]
            else:
                input_prompts_embed = None

            trace_dict["class_pairs"].append(class_pairs)
            trace_dict["class_id"].append(class_id)
            trace_dict["image_id"].append(image_id)
            trace_dict["interpolation_step"].append(interpolate_step)
            trace_dict["embed_len"].append(embed_len)
            trace_dict["pos_prompt_text"].append(pos_prompt)
            trace_dict["neg_prompt_text"].append(neg_prompt)
            trace_dict["input_file_path"].append(input_file)
            trace_dict["output_file_path"].append(output_file)
            trace_dict["input_prompts_embed"].append(input_prompts_embed)

    return trace_dict