add dropdown

app.py

@@ -1,70 +1,73 @@
 import gradio as gr
 import torch
-# from diffusers import AutoPipelineForText2Image
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, EulerDiscreteScheduler, UNet2DConditionModel
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from pathlib import Path
-import stone
-import requests
-import io
-import os
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 from PIL import Image
-import spaces
-
 import matplotlib.pyplot as plt
-import numpy as np
 from matplotlib.colors import hex2color
-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)
+import stone
+import os
 
-# Initial pipeline setup
-default_model = model_names[0]
-print(default_model)
-pipeline_text2image = DiffusionPipeline.from_pretrained(
-    default_model
-)
-pipeline_text2image = pipeline_text2image.to("cuda")
+# Define model initialization functions
+def load_model(model_name):
+    if model_name == "stabilityai/sdxl-turbo":
+        pipeline = DiffusionPipeline.from_pretrained(
+            model_name, 
+            torch_dtype=torch.float16, 
+            variant="fp16"
+        ).to("cuda")
+    elif model_name == "runwayml/stable-diffusion-v1-5":
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            model_name, 
+            torch_dtype=torch.float16
+        ).to("cuda")
+    elif model_name == "ByteDance/SDXL-Lightning":
+        base = "stabilityai/stable-diffusion-xl-base-1.0"
+        ckpt = "sdxl_lightning_4step_unet.safetensors"
+        unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+        unet.load_state_dict(load_file(hf_hub_download(model_name, ckpt), device="cuda"))
+        pipeline = StableDiffusionXLPipeline.from_pretrained(
+            base, 
+            unet=unet, 
+            torch_dtype=torch.float16, 
+            variant="fp16"
+        ).to("cuda")
+        pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config, timestep_spacing="trailing")
+    elif model_name == "segmind/SSD-1B":
+        pipeline = StableDiffusionXLPipeline.from_pretrained(
+            model_name, 
+            torch_dtype=torch.float16, 
+            use_safetensors=True, 
+            variant="fp16"
+        ).to("cuda")
+    else:
+        raise ValueError("Unknown model name")
+    return pipeline
+
+choices=[
+            "stabilityai/sdxl-turbo", 
+            "runwayml/stable-diffusion-v1-5", 
+            "ByteDance/SDXL-Lightning", 
+            "segmind/SSD-1B"
+        ]
+
+for model_name in choices:
+    load_model(model_name)
+
+# Initialize the default model
+default_model = "stabilityai/sdxl-turbo"
+
+pipeline_text2image = load_model(default_model)
 
-@spaces.GPU
 def getimgen(prompt):
-
-    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).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")
+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)
@@ -80,69 +83,37 @@ def genderfromcaption(caption):
 
 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
+    hex_codes = [code for code in hex_codes if code is not None]
     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]))
-
+        if i < len(sorted_hex_codes):
+            ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
     return fig
 
-@spaces.GPU
 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")
-
+    pipeline_text2image = load_model(model_name)
     foldername = "temp"
-    # Generate 10 images
-    images = [getimgen(prompt) for _ in range(10)]
-
     Path(foldername).mkdir(parents=True, exist_ok=True)
-
+    images = [getimgen(prompt) for _ in range(10)]
     genders = []
     skintones = []
-
     for image, i in zip(images, range(10)):
         prompt_prefix = "photo of a "
         caption = blip_caption_image(image, prefix=prompt_prefix)
@@ -153,26 +124,35 @@ def generate_images_plots(prompt, model_name):
             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 Text to Image Models") as demo:
-
+with gr.Blocks(title="Skin Tone and Gender bias in Text to Image Models") as 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)
+    model_dropdown = gr.Dropdown(
+        label="Choose a model", 
+        choices=[
+            "stabilityai/sdxl-turbo", 
+            "runwayml/stable-diffusion-v1-5", 
+            "ByteDance/SDXL-Lightning", 
+            "segmind/SSD-1B"
+        ], 
+        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")
+    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, model_dropdown], outputs=[gallery, skinplot, genplot])
 
-demo.launch(debug=True)
+demo.launch(debug=True)