app.py

import streamlit as st

import time, uuid
from datetime import datetime
import numpy as np
from PIL import Image

if 'session_id' not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
    st.session_state.n_gen = 0

# constants
HF_REPO_NAME_DIFFUSION = 'nostalgebraist/nostalgebraist-autoresponder-diffusion'
model_path_diffusion = 'nostalgebraist-autoresponder-diffusion'
timestep_respacing_sres1 = '20' # '90,60,60,20,20'
timestep_respacing_sres2 = '20' # '250'

DIFFUSION_DEFAULTS = dict(
    batch_size=1,
    n_samples=1,
    clf_free_guidance=True,
    clf_free_guidance_sres=False,
    guidance_scale=1,
    guidance_scale_sres=0,
    yield_intermediates=True
)

@st.experimental_singleton
def setup():
    import os, subprocess, sys
    if not os.path.exists('improved_diffusion'):
        os.system("git clone https://github.com/nostalgebraist/improved-diffusion.git")
    os.system("cd improved-diffusion && git fetch origin nbar-space && git checkout nbar-space && pip install -e .")
    os.system("pip install tokenizers x-transformers==0.22.0 axial-positional-embedding")
    os.system("pip install einops==0.3.2")
    sys.path.append("improved-diffusion")

    import improved_diffusion.pipeline
    from transformer_utils.util.tfm_utils import get_local_path_from_huggingface_cdn

    if not os.path.exists(model_path_diffusion):
        model_tar_name = 'model.tar'
        model_tar_path = get_local_path_from_huggingface_cdn(
            HF_REPO_NAME_DIFFUSION, model_tar_name
        )
        subprocess.run(f"tar -xf {model_tar_path} && rm {model_tar_path}", shell=True)

    checkpoint_path_sres1 = os.path.join(model_path_diffusion, "sres1.pt")
    config_path_sres1 = os.path.join(model_path_diffusion, "config_sres1.json")

    checkpoint_path_sres2 = os.path.join(model_path_diffusion, "sres2.pt")
    config_path_sres2 = os.path.join(model_path_diffusion, "config_sres2.json")

    # load
    sampling_model_sres1 = improved_diffusion.pipeline.SamplingModel.from_config(
        checkpoint_path=checkpoint_path_sres1,
        config_path=config_path_sres1,
        timestep_respacing=timestep_respacing_sres1
    )

    sampling_model_sres2 = improved_diffusion.pipeline.SamplingModel.from_config(
        checkpoint_path=checkpoint_path_sres2,
        config_path=config_path_sres2,
        timestep_respacing=timestep_respacing_sres2
    )

    pipeline = improved_diffusion.pipeline.SamplingPipeline(sampling_model_sres1, sampling_model_sres2)
    return pipeline


def now_str():
    return datetime.utcnow().strftime('%Y-%m-%d %H-%M-%S')


def log(msg, st_state):
    session_id = st_state.session_id if 'session_id' in st_state else None
    n_gen = st.session_state.n_gen if 'n_gen' in st_state else None

    print(f"{now_str()} {session_id} ({n_gen}th gen):\n\t{msg}\n")

def handler(text, ts1, ts2, gs1, st_state):
    pipeline = setup()

    data = {'text': text[:380], 'guidance_scale': gs1}
    args = {k: v for k, v in DIFFUSION_DEFAULTS.items()}
    args.update(data)

    log_data = {'ts1': ts2, 'ts2': ts2}
    log_data.update(args)

    log(repr(log_data), st_state)

    pipeline.base_model.set_timestep_respacing(str(ts1))
    pipeline.super_res_model.set_timestep_respacing(str(ts2))

    return pipeline.sample(**args)

FRESH = True

st.title('nostalgebraist-autoresponder image generation demo')

st.write("#### For a **much faster experience**, try the [Colab notebook](https://colab.research.google.com/drive/17BOTYmLv4fdurr8y5dcaGKy8JVY_A62a?usp=sharing) instead!")

st.write("A demo of the image models used in the tumblr bot [nostalgebraist-autoresponder](https://nostalgebraist-autoresponder.tumblr.com/).\n\nBy [nostalgebraist](https://nostalgebraist.tumblr.com/)")

st.write('##### What is this thing? How does it work?')

st.write("See [this post](https://nostalgebraist.tumblr.com/post/672300992964050944/franks-image-generation-model-explained) for an explanation.")

st.header('Prompt')

button_dril = st.button('Fill @dril tweet example text')

if FRESH and button_dril:
    st.session_state.fill_value = 'wint\nFollowing\n@dril\nthe wise man bowed his head solemnly and\nspoke: "theres actually zero difference\nbetween good & bad things. you imbecile.\nyou fucking moron'

if 'fill_value' in st.session_state:
    fill_value = st.session_state.fill_value
else:
    fill_value = ""
    st.session_state.fill_value = fill_value

text = st.text_area('Enter your text here (or leave blank for a textless image)', max_chars=380, height=230,
                    value=fill_value)

st.header('Settings')

st.write("The bot uses 250 base steps and 250 upsampling steps, with custom spacing (not available here) for the base part.\n\nSince this demo doesn't have a GPU, you'll probably want to use fewer than 250 steps unles you have a lot of patience.")

help_ts1 = "How long to run the base model. Larger values make the image more realistic / better. Smaller values are faster."
help_ts2 = "How long to run the upsampling model. Larger values sometimes make the big image crisper and more detailed. Smaller values are faster."
help_gs1 = "Guidance scale. Larger values make the image more likely to contain the text you wrote. If this is zero, the first part will be faster."

ts1 = st.slider('Steps (base)', min_value=5, max_value=500, value=50, step=5, help=help_ts1)
ts2 = st.slider('Steps (upsampling)', min_value=5, max_value=500, value=50, step=5, help=help_ts2)
gs1 = st.select_slider('Guidance scale (base)', [0.5*i for i in range(9)], value=1.0, help=help_gs1)

button_go = st.button('Generate')
button_stop = st.button('Stop')

st.write("During generation, the two images show different ways of looking at the same process.\n- The left image starts with 100% noise and gradually turns into 100 signal.\n- The right image shows the model's current 'guess' about what the left image will look like when all the noise has been removed.")

generating_marker = st.empty()
low_res = st.empty()
high_res = st.empty()

if button_go:
    st.session_state.n_gen = st.session_state.n_gen + 1
    with generating_marker.container():
        st.write('**Generating...**')
        st.write('**Prompt:**')
        st.write(repr(text))

    count_low_res, count_high_res = 0, 0
    times_low, times_high = [], []

    t = time.time()

    for s, xs in handler(text, ts1, ts2, gs1, st.session_state):
        s = Image.fromarray(s[0])
        xs = Image.fromarray(xs[0])

        t2 = time.time()
        delta = t2 - t
        t = t2

        is_high_res = s.size[0] == 256

        if is_high_res:
            target = high_res
            count_high_res += 1
            count = count_high_res
            total = ts2
            times_high.append(delta)
            times = times_high
            prefix = "Part 2 of 2 (upsampling)"
        else:
            target = low_res
            count_low_res += 1
            count = count_low_res
            total = ts1
            times_low.append(delta)
            times = times_low
            prefix = "Part 1 of 2 (base model)"

        rate = sum(times)/len(times)

        with target.container():
            st.image([s, xs])
            st.write(f'{prefix} | {count:02d} / {total} frames | {rate:.2f} seconds/frame')

        if button_stop:
            log('gen stopped', st.session_state)
            break

    with generating_marker.container():
        log('gen complete', st.session_state)
        st.write('')