File size: 5,914 Bytes
dbe3f9b
e584dab
dbe3f9b
 
b90fb1a
dbe3f9b
 
 
e584dab
dbe3f9b
 
 
e584dab
dbe3f9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e584dab
dbe3f9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e584dab
dbe3f9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e584dab
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
---
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.