Update app.py

app.py

@@ -6,7 +6,8 @@ from diffusers import (
     StableDiffusionXLPipeline,
     EulerDiscreteScheduler,
     UNet2DConditionModel,
-    StableDiffusion3Pipeline
+    StableDiffusion3Pipeline,
+    FluxPipeline
 )
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from pathlib import Path
@@ -21,11 +22,9 @@ import spaces
 
 access_token = os.getenv("AccessTokenSD3")
 
-
 from huggingface_hub import login
 login(token = access_token)
 
-
 # Define model initialization functions
 def load_model(model_name):
     if model_name == "stabilityai/sdxl-turbo":
@@ -65,6 +64,9 @@ def load_model(model_name):
             scheduler=scheduler, 
             torch_dtype=torch.float16
         ).to("cuda")
+    elif model_name == "black-forest-labs/FLUX.1-dev":
+        pipeline = FluxPipeline.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+        pipeline.enable_model_cpu_offload()
     else:
         raise ValueError("Unknown model name")
     return pipeline
@@ -76,16 +78,26 @@ pipeline_text2image = load_model(default_model)
 @spaces.GPU
 def getimgen(prompt, model_name):
     if model_name == "stabilityai/sdxl-turbo":
-        return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
+        return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2, height=512, width=512).images[0]
     elif model_name == "ByteDance/SDXL-Lightning":
-        return pipeline_text2image(prompt, num_inference_steps=4, guidance_scale=0).images[0]
+        return pipeline_text2image(prompt, num_inference_steps=4, guidance_scale=0, height=512, width=512).images[0]
     elif model_name == "segmind/SSD-1B":
         neg_prompt = "ugly, blurry, poor quality"
-        return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt).images[0]
+        return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt, height=512, width=512).images[0]
     elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
-        return pipeline_text2image(prompt=prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0).images[0]
+        return pipeline_text2image(prompt=prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0, height=512, width=512).images[0]
     elif model_name == "stabilityai/stable-diffusion-2":
-        return pipeline_text2image(prompt=prompt).images[0]
+        return pipeline_text2image(prompt=prompt, height=512, width=512).images[0]
+    elif model_name == "black-forest-labs/FLUX.1-dev":
+        return pipeline_text2image(
+            prompt,
+            height=512,
+            width=512,
+            guidance_scale=3.5,
+            num_inference_steps=50,
+            max_sequence_length=512,
+            generator=torch.Generator("cpu").manual_seed(0)
+        ).images[0]
 
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
@@ -130,30 +142,12 @@ def skintoneplot(hex_codes):
     return fig
 
 def age_detector(image):
-    """
-    A function that detects the age from an image.
-
-    Args:
-        image: The input image for age detection.
-
-    Returns:
-        str: The detected age label from the image.
-    """
     pipe = pipeline('image-classification', model="dima806/faces_age_detection", device=0)
     result = pipe(image)
     max_score_item = max(result, key=lambda item: item['score'])
     return max_score_item['label']
 
 def ageplot(agelist):
-    """
-    A function that plots age-related data based on the given list of age categories.
-
-    Args:
-        agelist (list): A list of age categories ("YOUNG", "MIDDLE", "OLD").
-
-    Returns:
-        fig: A matplotlib figure object representing the age plot.
-    """
     order = ["YOUNG", "MIDDLE", "OLD"]
     words = sorted(agelist, key=lambda x: order.index(x))
     colors = {"YOUNG": "skyblue", "MIDDLE": "royalblue", "OLD": "darkblue"}
@@ -166,39 +160,12 @@ def ageplot(agelist):
     return fig
 
 def is_nsfw(image):
-    """
-    A function that checks if the input image is not for all audiences (NFAA) by classifying it using 
-    an image classification pipeline and returning the label with the highest score.
-    
-    Args:
-        image: The input image to be classified.
-        
-    Returns:
-        str: The label of the NFAA category with the highest score.
-    """
     classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
     result = classifier(image)
     max_score_item = max(result, key=lambda item: item['score'])
     return max_score_item['label']
 
 def nsfwplot(nsfwlist):
-    """
-    Generates a plot of NFAA categories based on a list of NFAA labels.
-
-    Args:
-        nsfwlist (list): A list of NSFW labels ("normal" or "nsfw").
-
-    Returns:
-        fig: A matplotlib figure object representing the NSFW plot.
-
-    Raises:
-        None
-
-    This function takes a list of NFAA labels and generates a plot with a grid of 2 rows and 5 columns. 
-    Each label is sorted based on a predefined order and assigned a color. The plot is then created using matplotlib, 
-    with each cell representing an NFAA label. The color of each cell is determined by the corresponding label's color. 
-    The function returns the generated figure object.
-    """
     order = ["normal", "nsfw"]
     words = sorted(nsfwlist, key=lambda x: order.index(x))
     colors = {"normal": "mistyrose", "nsfw": "red"}
@@ -232,25 +199,21 @@ def generate_images_plots(prompt, model_name):
         except:
             skintones.append(None)
         genders.append(genderfromcaption(caption))
-        ages.append(age_detector(image))  # Call age_detector function
-        nsfws.append(is_nsfw(image))  # Call is_nsfw function
+        ages.append(age_detector(image))
+        nsfws.append(is_nsfw(image))
     return images, skintoneplot(skintones), genderplot(genders), ageplot(ages), nsfwplot(nsfws)
 
 with gr.Blocks(title="Demographic bias in Text-to-Image Generation Models") as demo:
     gr.Markdown("# Demographic bias in Text to Image Models")
     gr.Markdown('''
 In this demo, we explore the potential biases in text-to-image models by generating multiple images based on user prompts and analyzing the gender, skin tone, age, and potential sexual nature of the generated subjects. Here's how the analysis works:
-
 1. **Image Generation**: For each prompt, 10 images are generated using the selected model.
 2. **Gender Detection**: The [BLIP caption generator](https://huggingface.co/Salesforce/blip-image-captioning-large) is used to elicit gender markers by identifying words like "man," "boy," "woman," and "girl" in the captions.
 3. **Skin Tone Classification**: The [skin-tone-classifier library](https://github.com/ChenglongMa/SkinToneClassifier) is used to extract the skin tones of the generated subjects.
 4. **Age Detection**: The [Faces Age Detection model](https://huggingface.co/dima806/faces_age_detection) is used to identify the age of the generated subjects.
 5. **NFAA Detection**: The [Falconsai/nsfw_image_detection](https://huggingface.co/Falconsai/nsfw_image_detection) model is used to identify whether the generated images are NFAA (not for all audiences).
-
 #### Visualization
-
 We create visual grids to represent the data:
-
 - **Skin Tone Grids**: Skin tones are plotted as exact hex codes rather than using the Fitzpatrick scale, which can be [problematic and limiting for darker skin tones](https://arxiv.org/pdf/2309.05148).
 - **Gender Grids**: Light green denotes men, dark green denotes women, and grey denotes cases where the BLIP caption did not specify a binary gender.
 - **Age Grids**: Light blue denotes people between 18 and 30, blue denotes people between 30 and 50, and dark blue denotes people older than 50.
@@ -266,7 +229,8 @@ This demo provides an insightful look into how current text-to-image models hand
             "stabilityai/sdxl-turbo", 
             "ByteDance/SDXL-Lightning",
             "stabilityai/stable-diffusion-2",
-            "segmind/SSD-1B"
+            "segmind/SSD-1B",
+            "black-forest-labs/FLUX.1-dev"
         ], 
         value=default_model
     )