generate_images.py

# Copyright 2025 Adobe Research. All rights reserved.
# To view a copy of the license, visit LICENSE.md.

import os

import matplotlib.pyplot as plt
import pandas as pd
from PIL import Image
from tqdm import tqdm
import torch
from diffusers import (
    AutoPipelineForText2Image,
    EulerDiscreteScheduler,
    PixArtAlphaPipeline,
    StableDiffusionXLPipeline,
    UNet2DConditionModel,
    DiffusionPipeline,
    EDMDPMSolverMultistepScheduler,
    FluxPipeline,
    StableDiffusion3Pipeline,
    StableDiffusionPipeline,

)
from flask import Flask, jsonify, request
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file
import json

import weave
import wandb

from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
from diffusers import StableDiffusion3Pipeline
import torch

def initalize_model(model_name):
    if model_name == "sd3.5-large":
        model_id = "stabilityai/stable-diffusion-3.5-large"

        nf4_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        model_nf4 = SD3Transformer2DModel.from_pretrained(
            model_id,
            subfolder="transformer",
            quantization_config=nf4_config,
            torch_dtype=torch.bfloat16
        )

        pipeline = StableDiffusion3Pipeline.from_pretrained(
            model_id, 
            transformer=model_nf4,
            torch_dtype=torch.bfloat16
        )
        pipeline.enable_model_cpu_offload()
        params = {
            "num_inference_steps": 28,
            "guidance_scale": 4.5,
            "max_sequence_length": 512,
        }
        return pipeline, params
    elif model_name == "playground":
        pipe = DiffusionPipeline.from_pretrained(
            "playgroundai/playground-v2.5-1024px-aesthetic",
            torch_dtype=torch.float16,
            variant="fp16",
        ).to("cuda")
        params = {
            "num_inference_steps": 50,
            "guidance_scale": 3,
        }
        return pipe, params
    
    elif model_name == "dreamlike":
        pipe = DiffusionPipeline.from_pretrained(
            "dreamlike-art/dreamlike-photoreal-2.0",
            torch_dtype=torch.float16,
        ).to("cuda")
        params = {
            "num_inference_steps": 50,
            "guidance_scale": 3,
        }
        return pipe, params
    elif model_name == "flux":
        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, requires_safety_checker=False)
        pipe.enable_model_cpu_offload()
        params = {
            "num_inference_steps": 50,
            "guidance_scale": 3,
        }
        return pipe, params
    elif model_name == "sd-lightning":
        pipe = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16,
            variant="fp16",
        ).to("cuda")
        pipe.scheduler = EulerDiscreteScheduler.from_config(
            pipe.scheduler.config, timestep_spacing="trailing"
        )
        params = {
            "num_inference_steps": 16,
            "guidance_scale": 1,
        }
        return pipe, params
    elif model_name == "sd2":
        pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16)
        params = {
            "num_inference_steps": 50,
            "guidance_scale": 3,
        }
        return pipe, params
    elif model_name == "pixart":
        pipe = PixArtAlphaPipeline.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", torch_dtype=torch.float16)
        pipe.enable_model_cpu_offload()
        params = {
            "num_inference_steps": 50,
            "guidance_scale": 3,
        }
        return pipe, params
    else:
        raise ValueError(f"Model {model_name} not supported")

playground_pipe, playground_params = None, None
dreamlike_pipe, dreamlike_params = None, None
sd3_5_large_pipe, sd3_5_large_params = None, None
pixart_pipe, pixart_params = None, None
flux_pipe, flux_params = None, None
sd_lightning_pipe, sd_lightning_params = None, None
sd2_pipe, sd2_params = None, None

def generate_images(prompt, model, num_images=1, negative_prompt=None):
    """
    Generate images using various diffusion models.
    
    Args:
        prompt (str): The text prompt to generate images from
        model (str): The model to use ('playground', 'dreamlike', 'pixart', 'sd-lightning')
        num_images (int): Number of images to generate
        negative_prompt (str, optional): Negative prompt to guide generation
        
    Returns:
        list: List of PIL Image objects
    """
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    if model == "playground":
        pipe = playground_pipe
        # Optional: Use DPM++ 2M Karras scheduler for crisper fine details
        pipe.scheduler = EDMDPMSolverMultistepScheduler()
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **playground_params
        ).images
        
    elif model == "dreamlike":
        pipe = dreamlike_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            negative_prompt=negative_prompt,
            **dreamlike_params
        ).images
    elif model == "sd3.5-large":
        pipe = sd3_5_large_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **sd3_5_large_params
        ).images
        print(images)
        
    elif model == "flux":
        pipe = flux_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **flux_params
        ).images

    elif model == "sd-lightning":
        pipe = sd_lightning_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **sd_lightning_params
        ).images

    elif model == "pixart":
        pipe = pixart_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **pixart_params
        ).images
    elif model == "sd2":
        pipe = sd2_pipe
        images = pipe(
            prompt=prompt,
            num_images_per_prompt=num_images,
            **sd2_params
        ).images
    else:
        raise ValueError(f"Model {model} not supported")
    return images


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--prompts-file", type=str, default="data/prompts/bias.txt")
    parser.add_argument("--num-imgs-per-prompt", type=int, default=3)
    parser.add_argument("--prompt", type=str)
    parser.add_argument("--name", type=str)
    parser.add_argument("--models", type=str, nargs="+", default=["sd3.5-large", "sd2"])
    parser.add_argument(
        "--batch-size", type=int, default=3, help="Number of prompts to process in each batch"
    )
    args = parser.parse_args()

    for model in args.models:
        pipe, params = initalize_model(model)
        if model == "sd3.5-large":
            sd3_5_large_pipe, sd3_5_large_params = pipe, params
        elif model == "playground":
            playground_pipe, playground_params = pipe, params
        elif model == "dreamlike":
            dreamlike_pipe, dreamlike_params = pipe, params
        elif model == "flux":
            flux_pipe, flux_params = pipe, params
        elif model == "sd-lightning":
            sd_lightning_pipe, sd_lightning_params = pipe, params
        elif model == "pixart":
            pixart_pipe, pixart_params = pipe, params
        elif model == "sd2":
            sd2_pipe, sd2_params = pipe, params

    if args.prompt:
        prompts = [args.prompt]
        args.name = "SinglePrompt"
    else:
        if ".txt" in args.prompts_file:
            with open(args.prompts_file, "r") as file:
                prompts = [line.strip() for line in file.readlines()]
        else:
            data = pd.read_csv(args.prompts_file)
            prompts = list(set(data["Prompt"].unique())).sample(200)
        print(f"Loaded {len(prompts)} prompts from {args.prompts_file}")    

    
    wandb.init(
        project="Generate Text to Image",
        name=args.prompt.replace(" ", "_") if args.prompt else args.name,
        config=vars(args),
    )

    save_dir = f"data/{args.name}"
    # save_dir_pixart = f"{save_dir}/PixArt"
    # save_dir_sd_lightning = f"{save_dir}/SD-Lightning"

    # for dir in [save_dir, save_dir_pixart, save_dir_sd_lightning]:
    #     os.makedirs(dir, exist_ok=True)

    all_results = {
        "Prompt": [],
        "seed": [],
    }

    # Add "Model_" keys
    all_results.update({f"Model_{model}": [] for model in args.models})

    # Add "Image_" keys
    all_results.update({f"Image_{model}": [] for model in args.models})

    for i, prompt in enumerate(tqdm(prompts)):
        print(f"Generating {args.num_imgs_per_prompt} images for prompt: {prompt}")

        # Generate images for both models at once
        images = {model: generate_images(prompt, model, args.num_imgs_per_prompt) for model in args.models}
        fileid = (
            prompt.replace(" ", "_")
            .replace("?", "")
            .replace(":", "")
            .replace(",", "")
            .replace("!", "")
            .replace("/", "_")
            .replace(".", "")[:100]
        )
        for seed, (model, imgs) in enumerate(images.items()):
            # if seed == 0:
            #     all_results["Prompt"].append(prompt)
            #     all_results["seed"].append(seed)

            if not os.path.exists(f"{save_dir}/{model}"):
                os.makedirs(f"{save_dir}/{model}", exist_ok=True)

            print(f"Saving {len(imgs)} images for {model} to {save_dir}/{model}")

            for im_idx, im in enumerate(imgs):
                if seed == 0:
                    all_results["Prompt"].append(prompt)
                    all_results["seed"].append(seed)

                all_results[f"Model_{model}"].append(model)
                image_path = f"{save_dir}/{model}/{fileid}_{im_idx}.png"
                im.save(image_path)
                all_results[f"Image_{model}"].append(image_path)
        
        if args.num_imgs_per_prompt > 1 and len(args.models) > 1:
            # Create and log image grid for the prompt
            fig, axs = plt.subplots(len(args.models), min([5, args.num_imgs_per_prompt]), figsize=(20, 12))
            for i, (model, imgs) in enumerate(images.items()):
                for j, im in enumerate(imgs):
                    axs[i, j].imshow(im)
                    axs[i, j].axis("off")
            plt.tight_layout()
            wandb.log({prompt: wandb.Image(plt, caption=prompt)})
            plt.close(fig)

        # Save intermediate results
        for k in all_results.keys():
            print(k, len(all_results[k]))
        intermediate_df = pd.DataFrame(all_results)
        intermediate_df.to_json(
            f"{save_dir}/results_intermediate.json", orient="records", lines=True
        )

    # Save final results
    final_df = pd.DataFrame(all_results)
    final_df.to_json(f"{save_dir}/results_final.json", orient="records", lines=True)

    wandb_table = final_df.copy()
    for model in args.models:
        wandb_table[f"Image_{model}"] = [
            wandb.Image(Image.open(image_path)) for image_path in final_df[f"Image_{model}"]
        ]
    wandb.log({"final_results": wandb_table})

    # # Log final results to wandb
    # final_table = wandb.Table(dataframe=final_df)
    # wandb.log({"final_results_table": final_table})

    final_df["Prompt"].unique()
    new_df = pd.DataFrame(
        {
            "Prompt": pd.concat([final_df["Prompt"] for _ in args.models]),
            "subset": pd.concat([final_df["Prompt"] for _ in args.models]),
            "group_name": pd.concat([final_df[f"Model_{model}"] for model in args.models]),
            "path": pd.concat([final_df[f"Image_{model}"] for model in args.models]),
            "seed": pd.concat([final_df["seed"] for _ in args.models]),
        }
    )

    new_df = new_df.dropna()
    new_df.columns = ["Prompt", "subset", "group_name", "path", "seed"]
    new_df.to_csv(f"{save_dir}/results.csv", index=False)

wandb.finish()