app.py

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()