|
--- |
|
license: mit |
|
--- |
|
|
|
# worldly-v1: 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. |