worldly / README.md
WorldlyLabs's picture
Update README.md
b90fb1a verified
|
raw
history blame
5.91 kB
---
license: mit
---
# 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:
```bash
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:
```python
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
```python
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.