Implementation of the UI in Streamlit

[ZeroCool940711]

I'm trying to implement everything that is currently on the webui.py file but using streamlit, I'm not good with frontend stuff but streamlit makes some things easy to do and it has some features that I really like and think other people will also like so I decided to give it a try and see if I could contribute with those changes, I think I've asked a lot of feature requests and it's time I contribute in some way to this repo 😆

Here is a screenshot of the basic layout, I've tried to copy everything from the gradio layout but so far it's only the layout, I haven't done anything with the colors or anything else, so that's why it looks a lot basic compared to the gradio UI we currently have.

Note that the theme the UI is using is dark, that's not because of an extension I'm using or anything, streamlit has built-in themes and you can build your owns if you want, it has the basic dark and light themes so I just used the menu on the top right side and went to settings, that let you choose there between the light and dark theme, also, I'm using the Wide Mode to use all the available space, by default streamlit uses a centered layout, while its good I think the wide layout give more space and so I decided to use that.

In case you notice there is an option on the streamlit menu called "Rerun", this option lets you do manually a live reload so you can reload the ui with any new changes you made to the source code, as soon as you save your changes on the code streamlit will detect the changes and show you this message at the top, the "Always rerun" option and also an option on the settings called "Run on Save" let you automatically rerun the UI and load the new changes, as long as the code structure is done correctly and things are loaded outside the main thread reloading the app does not lose any information you previously introduced and is almost instant.

Another feature I like about streamlit is that it lets you change any element on the UI during execution or during loops, so, we could refresh the image on every step during the generation process, also create a progress bar that doesn't use a lot of resources.

Last but not least, streamlit lets you stop any generation or anything you're doing inside the UI at any time, in fact, any change made to the UI, be it moving the sliders, changing the prompt, the seed or even changing tabs will stop what is running on the background, there is also a Stop button that shows at the top when something is running and that stop button raises an exception that if caught we can do things like saving the images generated so far or even saving a video as we already have all the frames from the generated images in memory, that was something I did on a version of VQGAN+CLIP I was using before, I'm planning on using some of the stuff from there here if it is possible.

If anyone can help me with the layout or anything else I would really appreciate it, as I said, I'm not good with frontend stuff but I will try my best, I just hope it's enough.

[ZeroCool940711]

Hey guys, a little update on my progress for the implementation of the UI using Streamlit.
First thing first, I was not actually working on this until today because I wanted to wait until a more stable version of the main WebUI was released but I decided to start working on it today as I think it's going to take a while for things to slow down on the development of the UI seeing that it has become really popular lately, so, I've only working on it for half a day but I think so far I've managed to implement most of what is on the Text to Image tab, here is a GIF showing what's done so far:

Here are a couple of features it has as well as some improvements I made:

• Models are not loaded on during startup and instead are loaded when we hit the Generate button for the first time.
• Only the main model is kept in memory at all times after the first run (at least for now) and the other models like GFPGAN and RealESRGAN are only loaded if you select their corresponding checkboxes on the UI.
• If a model is already loaded in memory but you want to free some memory you can uncheck them on the UI and they will be removed from memory during the next generation or when you refresh the UI or do any action that re-runs the UI, later I will make it so they are removed from memory right away.
• Because of the previously mentioned points a lot of VRAM can be saved, so far I've seen it use around 5.2GB of RAM during generation while the gradio version uses around 7.8GB with the same options selected and all models in memory.
• The generation speed for some reason increased as well, around 75% even after upscaling and fixing the faces .-.
• Each slider and option on the UI has a small question button next to them, this will show a tooltip with extra information on what the option do and how to use them, this will reduce the amount of text and stuff on the UI as we can move any extra stuff there, you can probably see that some options have different names or smaller names compared to the gradio WebUI.
• Talking about gradio vs streamlit, if you need to kill the process the UI will reconnect automatically when you start the script again and in case of any changes on the code it will show a button at the top right to re-run the page which does a live reload, inserting any new changes on the page while keeping what didn't change, all in 1-3 seconds, so, messing with the UI or the code is a lot better as you do not need to be stopping the script all the time and waiting for minutes for the models to load, even when the UI is reloaded the models and some other information is still kept between sessions.
• I changed the folder names where the images are saved so they are now safe on both Linux and Windows and do not break with long prompts like before, I used the time each image was generated plus a unique ID.
• The information that was used to generate each image is still saved in a "YAML" if you select the option on the UI but just in case someone does not want the extra files I will also save the information as metadata on the image itself, that will be for tomorrow though.
• You can stop the generation at any time by either hitting the Stop button that will show at the top right when the generation is being done or you can also change something on the UI and it will stop the generation after the current operation finishes, if an image is being generated it will stop after that image is done or in the case a process like upscaling or fixing the faces is running it will then stop after that.
• For now the result image is shown after the generation like on gradio but I will make it soon so it refreshes while it's been generated so we can see how the process goes and if we don't like it then we can stop the generation, I think it's also possible to create a video to observe all the generation process but need to test it first to be sure.

Well, that's all for now, I'm going to sleep and tomorrow when I wake up I will continue to improve this, once I think it is in a usable state I will share the code so anyone who want to use it can use it, if you have some ideas of things that can be improved or things you want done differently just let me know and I will see what I can do.

[ZeroCool940711]

@hlky Could you give me some pointers on where to find the loop where the image is being generated inside webui.py, more specifically the place where each step is happening inside each iteration, I want to make it so on my streamlit version of the UI the image is updated as it's been generated, for now I manage to update it on the UI when it is saved by the save_sample() function which is a lot better than having to wait for all the batches to finish but I can't figure out exactly where each iteration and step is happening. I would really appreciate the help, my brain is currently completely fried and I can't figure that out by myself.

[altryne]

[codedealer]

webui 725 is the main loop if that's what you're asking

[ZeroCool940711]

Line 725 on my version or the original webui.py?

[codedealer]

https://github.com/hlky/stable-diffusion/blob/70bd4a4c21983292d3b396efaa026d6820869581/scripts/webui.py#L725

[hlky]

txt2img() and img2img() call process_images() with their func_init / init() and func_sample / sample() passed to process_images()

[ZeroCool940711]

Mmm, I'm looking for the line where it's told to generate each step, if you look at the console while there is an image been generated you can see something like this during the generation:

---

Iteration: 1/1
a red apple, highly detailed
1%|█ | 3/10 [00:01<00:00, 2.73it/s]

In this example you can see that there is 1 Iteration and that iteration has 10 steps or ddim_steps, I've already managed to push the image to the UI after each iteration when there are multiple iterations but what I want is to find the part where the ddim_steps are processed so I can display the image as it is being created on each step too but can't find that, it's probably a 1 line thing that is done so abstract that I'm missing it because I don't know what that line means, I've placed print statements all around the code trying to find that line and I've manage to print before the prompt is printed and after the bar is shown but not as the bar is being generated or the steps are being processed, that is what I'm looking for.

[hlky]

@ZeroCool940711 if you'd like I can create a branch for you to share your Streamlit implementation then people can help with the development. This is very impressive work, well done!

[ZeroCool940711]

Thanks, I think having it on a branch here would be a good idea. I will really appreciate any help I could get with this as there are areas where I'm lacking, for example I'm terrible with CSS and JS, there are things I would like to implement that I think I'm not going to be able to do alone or they would take me forever to do so, like for example the image that is being shown there in the middle, streamlit adds a full screen button so you can see the image a big bigger but the problem is that the image is shown on the left side instead of centered, that's a bit of a problem and it would be better if it was centered and allowed zooming in an out, those extra things could be really helpful, also I found an excellent image editor that is practically photoshop on the browser that would be perfect for the image to image tab but it requires JS knowledge and as I said I'm not good with it so I haven't been able to implement it yet. About the code, everything is on a single file and most of it is copied from the webui.py file, it has only an extra config file that I added for easy configuration as streamlit allows you to reload any config file without stopping the app, there were some people who wanted that for the WebUI in gradio but it was not implemented, my original idea was to have those options on the Settings menu that streamlit has but haven't implemented that yet.

[hlky]

I've created a streamlist branch on this repo, pull request your work :)

[ZeroCool940711]

Guys, I manage to finally found where the steps were generated on each iteration, it was inside each of the different samplers 🤦, so, for now I only manage to implement it for PLMS and DDIM, have to find a way to do it for the others as they are a bit different, what I did was to copy the classes from the respective samplers and add them to the webui_streamlit.py file so I could add a couple of lines to them to show the image on the UI after each step without affecting anything else like the webui.py which also uses it.

I don't know what other people think but seeing the image been generated in real time made me understand better what is going on behind the scenes, I also understood better what each option on the UI does as I was changing them and seeing the differences right away, for what I saw, right now we have little options to actually improve the quality of the result, there is actually one which is to increase the number of steps but that also increases the time it takes to generate each image, I do have something in mind but that would be probably for tomorrow or after I finish implementing this same feature on the other samplers, I want to see how the result changes with different samplers. Here is how it looks:

BTW, the performance was not affected that much, at least not in a noticeable way, I tested it multiple times to be sure with and without the images been rendered on each step and the speed was almost the same, here are some comparisons so you can see for yourself:

The first one is with the lines I added and the second one is as it was before, same with this other comparison, there is a small decrease in performance but it's not really something that you will notice, after all, what's 0.01-0.03 it/s compared to the ability to see the image been generated in real time.

I will be creating a PR soon with the changes in case someone wants to test it.

[altryne]

@ZeroCool940711 you won me over, seeing this happen in real time is going to be incredibly powerful feature.

👏

You said you aren't a frontend dev, I am, I will make this look nice, even tonight.

[altryne]

@ZeroCool940711 lmk your discord handle, I'll reach out with Qs

[inukai132]

@ZeroCool940711 that visualization is amazing

[TingTingin]

[ZeroCool940711]

@altryne I do suck at frontend stuff and would really appreciate if someone who knows their way around CSS and/or JS give me a hand with making the UI looks good, preferably making it look as similar to the gradio version as possible, remember that my idea is to make it easy for people to switch from one UI to the other and not to replace completely the gradio UI.
And my discord username is Tukirito#1748, add me there and let's talk more about how to improve this UI, it would be nice if this repo had a discord server and have there a channel for development so more people can give ideas and collaborate with development, nothing beats real time chat for sharing information, while getting comments here on GitHub does help a bit you get emails right away it's not as real time as discord.

[micwalk]

Does this have any relation to this other streamlit ui project? https://github.com/green-s/stable-diffusion-flow

[HenkDz]

Does this have any relation to this other streamlit ui project? https://github.com/green-s/stable-diffusion-flow

I don't think it does.

[codedealer]

Interesting that they're moving away to completely custom UI. Sooner or later all these large frameworks become more of a hindrance — good for fast prototyping but not for an actual product.

[TingTingin]

what's inpainting like in the new ui?

[ZeroCool940711]

Only the text to image tab is implemented right now on this new UI using streamlit but I do have some big plans for the image to image part like a better photo editor based on photoshop to have more options when editing an image.

[AK391]

great work using streamlit, however steramlit does not run in google colab afaik and gradio is planning on supporting interative outputs by end of the week @altryne @hlky

[ZeroCool940711]

You can always run the gradio based UI in colab, this UI in streamlit is an extra UI for those who want to use it, not one to replace the original.

[altryne]

@ZeroCool940711 I'm not sure the approach to support 2 UIs is going to work, even with 1 we have tons of bugs, merge conflicts, feature developments.

[ZeroCool940711]

You are right, 2 UIs might be too much, we then have to find a way to make streamlit work on colab, it probably can run just fine and its a matter of making small improvements to the code though.

[AK391]

gradio now supports iterative outputs, to use it you can try pip install gradio==3.2.1b0 @hlky @altryne

[codedealer]

Gradio 3.2x breaks some stuff in UI last time I checked, I wonder if they fixed it 🤔

[AK391]

@hlky @altryne have you been able to add iterative outputs to the gradio UI, let me know if you need any help here is the sample code

import numpy as np
import time
import os


def fake_diffusion(steps):
    for i in range(steps):
        time.sleep(1)
        yield np.random.random((200, 200, 3))
    time.sleep(1)
    yield "cheetah.jpg"

with gr.Blocks() as demo:
    slider = gr.Slider(1, 10, 3)
    button = gr.Button()
    image = gr.Image()
    button.click(fake_diffusion, slider, image)

demo.launch()

[codedealer]

Checked out the new interface. Feels nice, curious to look into how easy it is to develop for it. Looks like you can make custom components in react so basically anything should be possible to implement.

However it looks like my performance was hit:
DDIM - Streamlit with generation preview

DDIM - Streamlit without generation preview

DDIM - Gradio

Am I alone on this one?

[ZeroCool940711]

@codedealer sorry for taking so long to answer back, yeah, hot reloading makes things a lot easier, its a pain when you have to change one button and wait 5 minutes for the models to load then realize you need to change it again.

[ZeroCool940711]

BTW, here is a link for the discord server I'm setting up for hlky's repository and there is a channel for each version of the webui so devs can talk there about it and collaborate.
https://discord.gg/gyXNe4NySY

[altryne]

@codedealer sorry for taking so long to answer back, yeah, hot reloading makes things a lot easier, its a pain when you have to change one button and wait 5 minutes for the models to load then realize you need to change it again.

That's why we have the webui_playground

[codedealer]

That's why we have the webui_playground

And I appreciate that we do! But it's not the same as fully functioning hot reloading — you can't change the backend logic in playground.

[altryne]

And I appreciate that we do!
Agree, you can technically get gradio to reload by gradio scripts\webui_playgorund.py or even gradio webui.py but the latter will... reload the model EVERY time ugh

[JoshuaKimsey]

Whoa, this looks amazing! I really like the iterative display of the sampler output too, even though I know it likely will have negative impact on generation times. Would be curious to see where this goes.

[ZeroCool940711]

@JoshuaKimsey it's already on the main repo, check the webui_streamlit.py file, also, if you need help or want to help with the development check the discord server for the repo, here is the link

[AK391]

@JoshuaKimsey iterative outputs is also supported in gradio with a few lines of code and works inside google colab and kaggle notebook meaning you can use free gpu with a shareable link that can work anywhere on the web

import numpy as np
import time
import os


def fake_diffusion(steps):
    for i in range(steps):
        time.sleep(1)
        yield np.random.random((200, 200, 3))
    time.sleep(1)
    yield "cheetah.jpg"

with gr.Blocks() as demo:
    slider = gr.Slider(1, 10, 3)
    button = gr.Button()
    image = gr.Image()
    button.click(fake_diffusion, slider, image)

demo.launch()

[JoshuaKimsey]

@ZeroCool940711 Cool, I'll join the Discord, but there is no webui_streamlit.py in the main stable repo?

[JoshuaKimsey]

@AK391 Cool!

[ZeroCool940711]

mmm, it seems like it hasn't been made so it is synced with it, just get it from here as well as the webui_streamlit.yaml, just put the yaml file in configs/webui/webui_streamlit.yaml and the py file on the script folder, to run it install streamlit and python-slugify with pip like pip install streamlit python-slugify and then run the script with streamlit run scripts/webui_streamlit.py

[ZeroCool940711]

Hey guys, another update on the image generation preview part of the streamlit WebUI.
I finally manage to make use of the proper callback that the DDIM and PLMS samplers have, so now instead of the blocky and noisy image that we were getting before it now shows a blurry image but it gets more clear over time, this is way better than before and lets you see better the result you will get almost from the first step, here is how it looks like:

Hope you guys like it, I will be making a PR with the changes soon so keep an eye on the repo, also, the new changes will be publishes to the master branch from now on instead of been on the streamlit branch, the file called webui_streamlit.py will be probably on the scripts folder on the main repo too.

[ZeroCool940711]

Added a progress bar as well as some extra information to track the generation progress without having to check the console all the time.

[ZeroCool940711]

The Image to Image tab is now working at a basic level, there is no image editor for now, you can select an image using the "Upload Image" option on the left and the rest is just like the Text-to-Image tab, for now the upscaling model is disabled by default there as it has some problems I will be trying to fix tomorrow. Hope you guys like it.