playing around with model options

app copy.py

@@ -0,0 +1,149 @@
+import gradio as gr
+import torch
+from diffusers import AutoPipelineForText2Image
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from pathlib import Path
+import stone
+import requests
+import io
+import os
+from PIL import Image
+import spaces
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import hex2color
+
+
+pipeline_text2image = AutoPipelineForText2Image.from_pretrained(
+    "stabilityai/sdxl-turbo",
+    torch_dtype=torch.float16,
+    variant="fp16",
+)
+pipeline_text2image = pipeline_text2image.to("cuda")
+
+
+@spaces.GPU
+def getimgen(prompt):
+
+    return pipeline_text2image(
+        prompt=prompt,
+        guidance_scale=0.0,
+        num_inference_steps=2
+    ).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")
+
+
+@spaces.GPU
+def blip_caption_image(image, prefix):
+    inputs = blip_processor(image, prefix, return_tensors="pt").to("cuda", torch.float16)
+    out = blip_model.generate(**inputs)
+    return blip_processor.decode(out[0], skip_special_tokens=True)
+
+def genderfromcaption(caption):
+    cc = caption.split()
+    if "man" in cc or "boy" in cc:
+        return "Man"
+    elif "woman" in cc or "girl" in cc:
+        return "Woman"
+    return "Unsure"
+
+def genderplot(genlist):    
+    order = ["Man", "Woman", "Unsure"]
+
+    # Sort the list based on the order of keys
+    words = sorted(genlist, key=lambda x: order.index(x))
+
+    # Define colors for each category
+    colors = {"Man": "lightgreen", "Woman": "darkgreen", "Unsure": "lightgrey"}
+
+    # Map each word to its corresponding color
+    word_colors = [colors[word] for word in words]
+
+    # Plot the colors in a grid with reduced spacing
+    fig, axes = plt.subplots(2, 5, figsize=(5,5))
+
+    # Adjust spacing between subplots
+    plt.subplots_adjust(hspace=0.1, wspace=0.1)
+
+    for i, ax in enumerate(axes.flat):
+        ax.set_axis_off()
+        ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=word_colors[i]))
+
+    return fig
+
+def skintoneplot(hex_codes):
+    # Convert hex codes to RGB values
+    rgb_values = [hex2color(hex_code) for hex_code in hex_codes]
+
+    # Calculate luminance for each color
+    luminance_values = [0.299 * r + 0.587 * g + 0.114 * b for r, g, b in rgb_values]
+
+    # Sort hex codes based on luminance in descending order (dark to light)
+    sorted_hex_codes = [code for _, code in sorted(zip(luminance_values, hex_codes), reverse=True)]
+
+    # Plot the colors in a grid with reduced spacing
+    fig, axes = plt.subplots(2, 5, figsize=(5,5))
+
+    # Adjust spacing between subplots
+    plt.subplots_adjust(hspace=0.1, wspace=0.1)
+
+    for i, ax in enumerate(axes.flat):
+        ax.set_axis_off()
+        ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
+
+    return fig
+
+@spaces.GPU
+def generate_images_plots(prompt):
+    foldername = "temp"
+    # Generate 10 images
+    images = [getimgen(prompt) for _ in range(10)]
+
+    Path(foldername).mkdir(parents=True, exist_ok=True)
+
+    genders = []
+    skintones = []
+
+    for image, i in zip(images, range(10)):
+        prompt_prefix = "photo of a "
+        caption = blip_caption_image(image, prefix=prompt_prefix)
+        image.save(f"{foldername}/image_{i}.png")
+        try:
+            skintoneres = stone.process(f"{foldername}/image_{i}.png", return_report_image=False)
+            tone = skintoneres['faces'][0]['dominant_colors'][0]['color']
+            skintones.append(tone)
+        except:
+            skintones.append(None)
+
+        genders.append(genderfromcaption(caption))
+
+    print(genders, skintones)
+
+    return images, skintoneplot(skintones), genderplot(genders)
+
+
+with gr.Blocks(title = "Skin Tone and Gender bias in SDXL Demo - Inference API") as demo:
+
+    gr.Markdown("# Skin Tone and Gender bias in SDXL Demo")
+
+    prompt = gr.Textbox(label="Enter the Prompt")
+    gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", 
+                         columns=[5], rows=[2], object_fit="contain", height="auto")
+    btn = gr.Button("Generate images", scale=0)
+    with gr.Row(equal_height=True):
+        skinplot = gr.Plot(label="Skin Tone")
+        genplot = gr.Plot(label="Gender")
+
+
+    btn.click(generate_images_plots, inputs = prompt, outputs = [gallery, skinplot, genplot])
+
+
+
+demo.launch(debug=True)

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from diffusers import AutoPipelineForText2Image
+# from diffusers import AutoPipelineForText2Image
+from diffusers import DiffusionPipeline
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from pathlib import Path
 import stone
@@ -13,16 +14,41 @@ import spaces
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import hex2color
-
-
-pipeline_text2image = AutoPipelineForText2Image.from_pretrained(
-    "stabilityai/sdxl-turbo",
-    torch_dtype=torch.float16,
-    variant="fp16",
+from huggingface_hub import list_models
+
+# Fetch models from Hugging Face Hub
+models = list_models(task="text-to-image")
+## Step 1: Filter the models
+filtered_models = [model for model in models if model.library_name == "diffusers"]
+
+# Step 2: Sort the filtered models by downloads in descending order
+sorted_models = sorted(filtered_models, key=lambda x: x.downloads, reverse=True)
+
+# Step 3: Select the top 5 models with only one model per company
+top_models = []
+companies_seen = set()
+
+for model in sorted_models:
+    company_name = model.id.split('/')[0]  # Assuming the company name is the first part of the model id
+    if company_name not in companies_seen:
+        top_models.append(model)
+        companies_seen.add(company_name)
+    if len(top_models) == 5:
+        break
+
+# Get the ids of the top models
+model_names = [model.id for model in top_models]
+
+print(model_names)
+
+# Initial pipeline setup
+default_model = model_names[0]
+print(default_model)
+pipeline_text2image = DiffusionPipeline.from_pretrained(
+    default_model
 )
 pipeline_text2image = pipeline_text2image.to("cuda")
 
-
 @spaces.GPU
 def getimgen(prompt):
 
@@ -32,14 +58,12 @@ def getimgen(prompt):
         num_inference_steps=2
     ).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")
 
-
 @spaces.GPU
 def blip_caption_image(image, prefix):
     inputs = blip_processor(image, prefix, return_tensors="pt").to("cuda", torch.float16)
@@ -101,7 +125,15 @@ def skintoneplot(hex_codes):
     return fig
 
 @spaces.GPU
-def generate_images_plots(prompt):
+def generate_images_plots(prompt, model_name):
+    print(model_name)
+    # Update the pipeline to use the selected model
+    global pipeline_text2image
+    pipeline_text2image = DiffusionPipeline.from_pretrained(
+        model_name
+    )
+    pipeline_text2image = pipeline_text2image.to("cuda")
+
     foldername = "temp"
     # Generate 10 images
     images = [getimgen(prompt) for _ in range(10)]
@@ -128,11 +160,11 @@ def generate_images_plots(prompt):
 
     return images, skintoneplot(skintones), genderplot(genders)
 
+with gr.Blocks(title = "Skin Tone and Gender bias in Text to Image Models") as demo:
 
-with gr.Blocks(title = "Skin Tone and Gender bias in SDXL Demo - Inference API") as demo:
-
-    gr.Markdown("# Skin Tone and Gender bias in SDXL Demo")
+    gr.Markdown("# Skin Tone and Gender bias in Text to Image Models")
 
+    model_dropdown = gr.Dropdown(label="Choose a model", choices=model_names, value=default_model)
     prompt = gr.Textbox(label="Enter the Prompt")
     gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", 
                          columns=[5], rows=[2], object_fit="contain", height="auto")
@@ -141,9 +173,6 @@ with gr.Blocks(title = "Skin Tone and Gender bias in SDXL Demo - Inference API")
         skinplot = gr.Plot(label="Skin Tone")
         genplot = gr.Plot(label="Gender")
 
+    btn.click(generate_images_plots, inputs=[prompt, model_dropdown], outputs=[gallery, skinplot, genplot])
 
-    btn.click(generate_images_plots, inputs = prompt, outputs = [gallery, skinplot, genplot])
-
-
-
-demo.launch(debug=True)
+demo.launch(debug=True)

test.ipynb

@@ -0,0 +1,277 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffusers import AutoPipelineForText2Image\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "vae/diffusion_pytorch_model.safetensors not found\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1822a5456c3244b6b5831817d6d0ebbc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 15 files:   0%|          | 0/15 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6232a45e054e4439a822345b4744b0b8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer/special_tokens_map.json:   0%|          | 0.00/472 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "da77db92e94e482b83b679b698630be0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "(…)ature_extractor/preprocessor_config.json:   0%|          | 0.00/342 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1f73d5e96b8b4c218433d8c9ea66b4b1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "safety_checker/config.json:   0%|          | 0.00/4.80k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4eb1d6d2c98c46e6940129236928b9d8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "text_encoder/config.json:   0%|          | 0.00/589 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8455fd4bd49e462e94a15614968edeeb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer/merges.txt:   0%|          | 0.00/525k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "85c9b8ddc3084051b9de01a1cc4b346f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "scheduler/scheduler_config.json:   0%|          | 0.00/284 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f01d373de73e49198ed9a509ff884d83",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "unet/config.json:   0%|          | 0.00/789 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2dad6d09e0b944cb91e5e68d374ad283",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer/vocab.json:   0%|          | 0.00/1.06M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "69c41904a7dd439296a3a56c0a37acda",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "vae/config.json:   0%|          | 0.00/592 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a3c8d76377f9445e81d73336d2b70535",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "safety_checker/pytorch_model.bin:   0%|          | 0.00/608M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "67a6fd052c734c69950ce38c7503821a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "text_encoder/pytorch_model.bin:   0%|          | 0.00/246M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e7de26e3010f484bb698f62dfae0a255",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer/tokenizer_config.json:   0%|          | 0.00/805 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d4080d04c64642a09514fd7570ba0cde",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "unet/diffusion_pytorch_model.bin:   0%|          | 0.00/1.72G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "402fd6d5c4054bdfae85f085c538ebac",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "vae/diffusion_pytorch_model.bin:   0%|          | 0.00/167M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# model = \"CompVis/ldm-text2im-large-256\"\n",
+    "model = \"sd-dreambooth-library/colorful-ball\"\n",
+    "# model = \"stabilityai/sdxl-turbo\"\n",
+    "\n",
+    "pipeline_text2image = AutoPipelineForText2Image.from_pretrained(\n",
+    "    model,\n",
+    "    torch_dtype=torch.float16,\n",
+    ")\n",
+    "pipeline_text2image = pipeline_text2image.to(\"cuda\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gradio",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}