worldly: Bias Mitigation Script for Image Generation

Overview

worldly is a bias mitigation script designed to modify prompts before sending them to an image generation model. It introduces diverse ethnicities and other demographic characteristics into prompts that contain vague references to "people," "person," or related terms, helping ensure more equitable representation in generated images. This version is specifically demonstrated with the FluxPipeline model, but it can be used with any image generation model that accepts text-based prompts.

Purpose

The goal of worldly is to mitigate bias in AI-generated imagery by diversifying the representations of people in prompts. This script dynamically modifies prompts by injecting randomly selected ethnicities or demographic details, ensuring equal chances of different ethnicities being represented in the generated images.

How It Works

• The script targets terms like "person," "people," "man," "woman," "child," "boy," "girl," and their plurals.
• It replaces these terms with a randomly selected ethnicity or demographic detail based on a detailed list of major ethnic and racial groups.
• The modified prompt is then passed to the image generation model to create more diverse and inclusive images.

Installation and Setup

Requirements

Make sure you have the following Python libraries installed:

pip install torch diffusers huggingface_hub Pillow

How to Use

1. Download the Script

You can download and integrate the worldly script into your image generation pipeline. Use the huggingface_hub library to fetch the script:

from huggingface_hub import hf_hub_download
import importlib.util

repo_id = "WorldlyLabs/worldly"
filename = "worldly-v1.py"
script_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load the script dynamically
spec = importlib.util.spec_from_file_location("worldly-v1", script_path)
worldly_v1 = importlib.util.module_from_spec(spec)
spec.loader.exec_module(worldly_v1)

2. Integrating into Image Generation

Once downloaded, you can use worldly to modify prompts before image generation. Below is an example of how to integrate the script into an image generation pipeline using the FluxPipeline model from Diffusers.

Example Script

import os
import torch
import gc
import logging
from huggingface_hub import hf_hub_download
import importlib.util
from diffusers import FluxPipeline

# Set up logging to print to console
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

# Download the worldly-v1 script from Hugging Face
repo_id = "WorldlyLabs/worldly"
filename = "worldly-v1.py"
script_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load the worldly-v1 script dynamically
spec = importlib.util.spec_from_file_location("worldly-v1", script_path)
worldly_v1 = importlib.util.module_from_spec(spec)
spec.loader.exec_module(worldly_v1)

# List of example prompts
prompts = {
    "Sample 1": "A person standing in a forest, looking at the sky.",
    "Sample 2": "A group of people walking in a desert, wearing traditional clothing.",
    "Sample 3": "A child holding a kite on a beach, with waves crashing nearby."
}

# Apply the worldly-v1 script to all prompts before image generation
def modify_all_prompts(prompts):
    modified_prompts = {}
    for prompt_name, prompt in prompts.items():
        modified_prompt = worldly_v1.modify_prompt(prompt)
        logging.info(f"Original prompt for {prompt_name}: {prompt}")
        logging.info(f"Modified prompt for {prompt_name}: {modified_prompt}")
        modified_prompts[prompt_name] = modified_prompt
    return modified_prompts

# Function to generate images for each prompt
def generate_images():
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        weight_dtype = torch.bfloat16

        # Load the FLUX pipeline
        print("Loading the FLUX pipeline...")
        pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            torch_dtype=weight_dtype,
        )
        
        pipe.to(device)
        pipe.enable_model_cpu_offload()

        # Modify all prompts before generating images
        modified_prompts = modify_all_prompts(prompts)

        # Create a folder for the generated images
        output_folder = "./generated_images"
        os.makedirs(output_folder, exist_ok=True)

        # Generate and save each image based on the modified prompt
        for prompt_name, prompt in modified_prompts.items():
            print(f"Generating image for {prompt_name}")

            # Generate the image
            prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt)

            image = pipe(
                prompt_embeds=prompt_embeds,
                pooled_prompt_embeds=pooled_prompt_embeds,
                guidance_scale=3.5,
                output_type="pil",
                num_inference_steps=80,
                max_sequence_length=256,
                generator=torch.Generator("cpu").manual_seed(0),
                height=1920,
                width=1080
            ).images[0]

            # Save the generated image with the prompt name
            output_path = os.path.join(output_folder, f"{prompt_name}.png")
            image.save(output_path)
            print(f"Image for {prompt_name} saved at: {output_path}")

        # Clear CPU cache
        gc.collect()

    except Exception as e:
        logging.error(f"Error during image generation: {str(e)}")

# Main execution block
if __name__ == "__main__":
    generate_images()

License

The worldly script is licensed under the MIT License. You are free to use, modify, and distribute this script, as long as the original copyright and permission notice is retained.