Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import numpy as np | |
| from PIL import Image | |
| # 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=False, | |
| clf_free_guidance_sres=False, | |
| guidance_scale=1, | |
| guidance_scale_sres=0, | |
| yield_intermediates=True | |
| ) | |
| 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 | |
| pipeline = setup() | |
| def handler(text, ts1, ts2, gs1): | |
| # a = np.random.randint(0, 255, (128, 128, 3)).astype(np.uint8) | |
| data = {'text': text[:380], 'guidance_scale': gs1} | |
| args = {k: v for k, v in DIFFUSION_DEFAULTS.items()} | |
| args.update(data) | |
| print(f"running: {args}") | |
| pipeline.base_model.set_timestep_respacing(str(ts1)) | |
| pipeline.super_res_model.set_timestep_respacing(str(ts2)) | |
| return pipeline.sample(**args) | |
| text = st.text_area('asdf') | |
| if st.button('rweerew'): | |
| low_res = st.empty() | |
| high_res = st.empty() | |
| for s, xs in handler(text, 20, 20, 0): | |
| print((type(s), s.shape)) | |
| print((type(xs), s.shape)) | |
| s = Image.fromarray(s[0]) | |
| xs = Image.fromarray(xs[0]) | |
| target = low_res if s.size[0] < 256 else high_res | |
| with target.container(): | |
| st.image([s, xs]) | |
| # x = st.slider('Select a value') | |
| # st.write(x, 'squared is', x * x) | |