aroffe/comparing-diffusion-models

1import gradio as gr2import os3 4from huggingface_hub import HfApi, ModelFilter, list_liked_repos, SpaceHardware5from diffusers import DiffusionPipeline6import torch7 8def gpu_enabled() -> bool:9    # If cloned, fill in SPACE_ID with your own space10    SPACE_ID = "aroffe/comparing-diffusion-models"11    runtime = api.get_space_runtime(repo_id=SPACE_ID)12    return (runtime.hardware != SpaceHardware.CPU_BASIC and runtime.hardware != SpaceHardware.CPU_UPGRADE)13 14    15def image_mod(prompt: str, model: str, image_0: gr.Image, image_1: gr.Image) -> list[gr.Image]:16    gpu_enabled()17    images = [image_0, image_1]18    for i, diffusion_model in enumerate(model):19        pipeline = DiffusionPipeline.from_pretrained(20            pretrained_model_name_or_path=diffusion_model,21            torch_dtype=torch.float16 if gpu_enabled() else torch.float32,22            use_safetensors=True,23            device_map="auto"24        )25        images[i] = pipeline(prompt).images[0]26    return images27 28 29def get_text_to_image_models():30    api = HfApi()31 32    # Out app utilizers Diffusers and Safetensors33    model_filter = ModelFilter(34        library=["diffusers" , "safetensors"],35        tags=["text-to-image"]36    )37 38    return api.list_models(filter=model_filter)39 40 41def get_liked_models(username: str) -> list:42    user_liked_models = list_liked_repos(username)43    text_to_image_models = get_text_to_image_models()44    text_to_image_models_ids = [model.id for model in list(text_to_image_models)]45 46    # find intersection between user liked models and models pulled from huggingface47    user_liked_text_to_image_models = [i for i in user_liked_models.models if i in text_to_image_models_ids]48    return user_liked_text_to_image_models49 50 51def get_dropdown_for_huggingface(profile: gr.OAuthProfile):52    return gr.Dropdown(53        choices=get_liked_models(profile.username),54        multiselect=True,55        max_choices=2,56        label="Select a Model from Hugginface",57        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.",58        every=60,59        visible=True60    )61 62 63def change_model(choice): 64    images = [gr.Image(visible=False), gr.Image(visible=False)]65    for i, model in enumerate(choice):66        images[i] = gr.Image(label=model, visible=True, interactive=False)67    68    return images69 70 71with gr.Blocks() as demo:72    gr.Markdown(73    """74    # Comparing Stable Diffusion Models75 76    ## Motivation77 78    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).79    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!80 81    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).82    83    I'd love to hear your feedback!84 85    ## How to use this app86 87    This app uses a [DiffusionPipeline](https://huggingface.co/docs/diffusers/main/en/api/diffusion_pipeline) to generate images.88 89    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.90 91    ## Speed up Inference92    93    For faster inference, you can duplicate the space and upgrade the hardware resources!94    """)95    # Make the buttons the same size for aesthetics96    gr.DuplicateButton()97    gr.LoginButton()98    dropdown = gr.Dropdown(visible=False)99    demo.load(get_dropdown_for_huggingface, inputs=None, outputs=dropdown)100    101    textbox = gr.Textbox(label="Prompt for Image Generation")102    103    with gr.Row() as row:104        image_0 = gr.Image(visible=False)105        image_1 = gr.Image(visible=False)106        dropdown.change(107            fn=change_model,108            inputs=dropdown,109            outputs=[image_0, image_1]110        )111    112    textbox.submit(113        image_mod, [textbox, dropdown, image_0, image_1], [image_0, image_1]114    )115 116demo.launch()117