aroffe's picture
Update app.py
096fcbf verified
import gradio as gr
import os
from huggingface_hub import HfApi, ModelFilter, list_liked_repos, SpaceHardware
from diffusers import DiffusionPipeline
import torch
def gpu_enabled() -> bool:
# If cloned, fill in SPACE_ID with your own space
SPACE_ID = "aroffe/comparing-diffusion-models"
runtime = api.get_space_runtime(repo_id=SPACE_ID)
return (runtime.hardware != SpaceHardware.CPU_BASIC and runtime.hardware != SpaceHardware.CPU_UPGRADE)
def image_mod(prompt: str, model: str, image_0: gr.Image, image_1: gr.Image) -> list[gr.Image]:
gpu_enabled()
images = [image_0, image_1]
for i, diffusion_model in enumerate(model):
pipeline = DiffusionPipeline.from_pretrained(
pretrained_model_name_or_path=diffusion_model,
torch_dtype=torch.float16 if gpu_enabled() else torch.float32,
use_safetensors=True,
device_map="auto"
)
images[i] = pipeline(prompt).images[0]
return images
def get_text_to_image_models():
api = HfApi()
# Out app utilizers Diffusers and Safetensors
model_filter = ModelFilter(
library=["diffusers" , "safetensors"],
tags=["text-to-image"]
)
return api.list_models(filter=model_filter)
def get_liked_models(username: str) -> list:
user_liked_models = list_liked_repos(username)
text_to_image_models = get_text_to_image_models()
text_to_image_models_ids = [model.id for model in list(text_to_image_models)]
# find intersection between user liked models and models pulled from huggingface
user_liked_text_to_image_models = [i for i in user_liked_models.models if i in text_to_image_models_ids]
return user_liked_text_to_image_models
def get_dropdown_for_huggingface(profile: gr.OAuthProfile):
return gr.Dropdown(
choices=get_liked_models(profile.username),
multiselect=True,
max_choices=2,
label="Select a Model from Hugginface",
info="Select up to two models to compare. If you want to see more models populated in the dropdown, just like a model on huggingface! The dropdown is refreshed every minute.",
every=60,
visible=True
)
def change_model(choice):
images = [gr.Image(visible=False), gr.Image(visible=False)]
for i, model in enumerate(choice):
images[i] = gr.Image(label=model, visible=True, interactive=False)
return images
with gr.Blocks() as demo:
gr.Markdown(
"""
# Comparing Stable Diffusion Models
## Motivation
Much of the motivation of this project was to compare open-source text-to-image models (I'd like to expand this beyond open-source as well).
In the process of building this, I thought of the use case where a developer is developing there own model -- or maybe fine-tuning an existing one -- and wants to compare the output of said model against a published one. A side-by-side tool can be adapted to that!
I see this space evolving in the future to include a comprehensive list of popular text-to-image model, saving the user from going into Hugging Face and liking models to see them in this app. (Anything but providing models by free text was my goal).
I'd love to hear your feedback!
## How to use this app
This app uses a [DiffusionPipeline](https://huggingface.co/docs/diffusers/main/en/api/diffusion_pipeline) to generate images.
In the dropdown you will see models that you have liked on HuggingFace. Only models that you've likes and fit the criteria for this app will display in the dropdown. You can navigate [here](https://huggingface.co/models?pipeline_tag=text-to-image&library=diffusers,safetensors&sort=trending) to get models that will work with this app. Go to the model card and like it; then upon refreshing this app you will see this model in the dropdown.
## Speed up Inference
For faster inference, you can duplicate the space and upgrade the hardware resources!
""")
# Make the buttons the same size for aesthetics
gr.DuplicateButton()
gr.LoginButton()
dropdown = gr.Dropdown(visible=False)
demo.load(get_dropdown_for_huggingface, inputs=None, outputs=dropdown)
textbox = gr.Textbox(label="Prompt for Image Generation")
with gr.Row() as row:
image_0 = gr.Image(visible=False)
image_1 = gr.Image(visible=False)
dropdown.change(
fn=change_model,
inputs=dropdown,
outputs=[image_0, image_1]
)
textbox.submit(
image_mod, [textbox, dropdown, image_0, image_1], [image_0, image_1]
)
demo.launch()