Spaces:
Runtime error
Runtime error
| from functools import partial | |
| from random import randint | |
| import gradio as gr | |
| import torch | |
| from tqdm import tqdm | |
| from NestedPipeline import NestedStableDiffusionPipeline | |
| from NestedScheduler import NestedScheduler | |
| def run(prompt, outer, inner, random_seed, pipe): | |
| seed = 24 if not random_seed else randint(0, 10000) | |
| generator = torch.Generator(device).manual_seed(seed) | |
| outer_diffusion = tqdm(range(outer), desc="Outer Diffusion") | |
| inner_diffusion = tqdm(range(inner), desc="Inner Diffusion") | |
| cur = [0, 0] | |
| for i, j, im in pipe(prompt, num_inference_steps=outer, num_inner_steps=inner, generator=generator): | |
| if cur[-1] != j: | |
| inner_diffusion.update() | |
| cur[-1] = j | |
| if cur[0] != i and i != outer: | |
| cur[0] = i | |
| outer_diffusion.update() | |
| cur[-1] = 0 | |
| inner_diffusion = tqdm(range(inner), desc="Inner Diffusion") | |
| elif cur[0] != i: | |
| outer_diffusion.update() | |
| monospace_s, monospace_e = "<p style=\"font-family:'Lucida Console', monospace\">", "</p>" | |
| yield f"{monospace_s}{outer_diffusion.__str__().replace(' ', ' ')}{monospace_e} \n {monospace_s}{inner_diffusion.__str__().replace(' ', ' ')}{monospace_e}", im[0] | |
| if __name__ == "__main__": | |
| scheduler = NestedScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", | |
| prediction_type='sample', clip_sample=False, set_alpha_to_one=False) | |
| fp16 = False | |
| if fp16: | |
| pipe = NestedStableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", revision="fp16", | |
| torch_dtype=torch.float16, scheduler=scheduler) | |
| else: | |
| pipe = NestedStableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", scheduler=scheduler) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| pipe.to(device) | |
| interface = partial(run, pipe=pipe) | |
| demo = gr.Interface( | |
| fn=interface, | |
| title="Nested Diffusion", | |
| description="<h3 style=\"text-align: center;\">Anytime text-to-image generation with Stable Diffusion v1.5</h3>\n<p style=\"text-align: center;\"><b>Help: </b>Type the desired prompt in the prompt box, and adjust the number of outer and inner steps to use. Using more steps takes more time, but should create a better image.<br>For more information on Nested Diffuion: <a href=\"https://github.com/noamelata/NestedDiffusion\">Github</a>, <a href=\"https://arxiv.org/abs/2305.19066\">arXiv</a></p>", | |
| inputs=[gr.Textbox(value="a photograph of a nest with a blue egg inside", label="Prompt"), | |
| gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Outer Steps"), | |
| gr.Slider(minimum=5, maximum=50, value=10, step=1, label="Inner Steps"), | |
| gr.Checkbox(label="Random Seed")], | |
| outputs=[gr.HTML(), gr.Image(shape=[512, 512], elem_id="output_image").style(width=512, height=512)], | |
| allow_flagging="never", | |
| thumbnail="figures/Nested_Egg.png" | |
| ) | |
| demo.queue() | |
| demo.launch() | |