Update app.py

app.py

@@ -76,7 +76,10 @@ def show_images(batch: th.Tensor):
     display(Image.fromarray(reshaped.numpy()))
 
 
-def compose_language_descriptions(prompt, guidance_scale):
+def compose_language_descriptions(prompt, guidance_scale, steps):
+    options['timestep_respacing'] = str(steps)
+    _, diffusion = create_model_and_diffusion(**options)
+
     # @markdown `prompt`: when composing  multiple sentences, using `|` as the delimiter.
     prompts = [x.strip() for x in prompt.split('|')]
 
@@ -240,12 +243,15 @@ clevr_model.load_state_dict(th.load(download_model('clevr_pos'), device))
 print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))
 
 
-def compose_clevr_objects(prompt, guidance_scale):
+def compose_clevr_objects(prompt, guidance_scale, steps):
     coordinates = [[float(x.split(',')[0].strip()), float(x.split(',')[1].strip())]
                    for x in prompt.split('|')]
     coordinates += [[-1, -1]]  # add unconditional score label
     batch_size = 1
 
+    clevr_options['timestep_respacing'] = str(int(steps))
+    _, clevr_diffusion = create_model_and_diffusion_for_clevr(**clevr_options)
+
     def model_fn(x_t, ts, **kwargs):
         half = x_t[:1]
         combined = th.cat([half] * kwargs['y'].size(0), dim=0)
@@ -281,24 +287,23 @@ def compose_clevr_objects(prompt, guidance_scale):
     out_img = (out_img + 1) / 2
     out_img = (out_img.detach().cpu() * 255.).to(th.uint8)
     out_img = out_img.numpy()
-    Image.fromarray(out_img).convert('RGB').save('test.png')
 
     return out_img
 
 
-def stable_diffusion_compose(prompt, scale):
+def stable_diffusion_compose(prompt, scale, steps):
     with autocast('cpu' if not has_cuda else 'cuda'):
-        image = pipe(prompt, guidance_scale=scale)["sample"][0]
+        image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps)["sample"][0]
         return image
 
 
-def compose(prompt, version, guidance_scale):
+def compose(prompt, version, guidance_scale, steps):
     if version == 'GLIDE':
-        return compose_language_descriptions(prompt, guidance_scale)
+        return compose_language_descriptions(prompt, guidance_scale, steps)
     elif version == 'Stable_Diffusion_1v_4':
-        return stable_diffusion_compose(prompt, guidance_scale)
+        return stable_diffusion_compose(prompt, guidance_scale, steps)
     else:
-        return compose_clevr_objects(prompt, guidance_scale)
+        return compose_clevr_objects(prompt, guidance_scale, steps)
 
 
 examples_1 = 'a camel | a forest'
@@ -309,20 +314,28 @@ examples_5 = 'a white church on a hill | birds flying around the church'
 examples_6 = 'a boat in a desert | a pink sky'
 examples_7 = 'mountains in the background |  a blue sky | cows on a pasture'
 examples = [
-            [examples_7, 'Stable_Diffusion_1v_4', 10],
-            [examples_4, 'Stable_Diffusion_1v_4', 10],
-            [examples_5, 'Stable_Diffusion_1v_4', 10],
-            [examples_6, 'Stable_Diffusion_1v_4', 10],
-            [examples_1, 'GLIDE', 10],
-            [examples_2, 'GLIDE', 10],
-            [examples_3, 'CLEVR Objects', 10]]
+            [examples_7, 'Stable_Diffusion_1v_4', 10, 50],
+            [examples_4, 'Stable_Diffusion_1v_4', 10, 50],
+            [examples_5, 'Stable_Diffusion_1v_4', 10, 50],
+            [examples_6, 'Stable_Diffusion_1v_4', 10, 50],
+            [examples_1, 'GLIDE', 10, 100],
+            [examples_2, 'GLIDE', 10, 100],
+            [examples_3, 'CLEVR Objects', 10, 100]
+]
 
 import gradio as gr
 
 title = 'Compositional Visual Generation with Composable Diffusion Models'
-description = '<p>Demo for Composable Diffusion<ul><li>~30s per GLIDE/Stable-Diffusion example</li><li>~10s per CLEVR Object example</li>(<b>Note</b>: time is varied depending on what gpu is used.)</ul></p><p>See more information from our <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/">Project Page</a>.</p><ul><li>One version is based on the released <a href="https://github.com/openai/glide-text2im">GLIDE</a> and <a href="https://github.com/CompVis/stable-diffusion/">Stable Diffusion</a> for composing natural language description.</li><li>Another is based on our pre-trained CLEVR Object Model for composing objects. <br>(<b>Note</b>: We recommend using <b><i>x</i></b> in range <b><i>[0.1, 0.9]</i></b> and <b><i>y</i></b> in range <b><i>[0.25, 0.7]</i></b>, since the training dataset labels are in given ranges.)</li></ul><p>When composing  multiple sentences, use `|` as the delimiter, see given examples below.</p>'
-
-iface = gr.Interface(compose, inputs=["text", gr.Radio(['Stable_Diffusion_1v_4', 'GLIDE', 'CLEVR Objects'], type="value", label='version'), gr.Slider(2, 20)], outputs='image',
+description = '<p>Demo for Composable Diffusion<ul><li>~30s per GLIDE/Stable-Diffusion example</li><li>~10s per CLEVR Object example</li>(<b>Note</b>: time is varied depending on what gpu is used.)</ul></p><p>See more information from our <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/">Project Page</a>.</p><ul><li>One version is based on the released <a href="https://github.com/openai/glide-text2im">GLIDE</a> and <a href="https://github.com/CompVis/stable-diffusion/">Stable Diffusion</a> for composing natural language description.</li><li>Another is based on our pre-trained CLEVR Object Model for composing objects. <br>(<b>Note</b>: We recommend using <b><i>x</i></b> in range <b><i>[0.1, 0.9]</i></b> and <b><i>y</i></b> in range <b><i>[0.25, 0.7]</i></b>, since the training dataset labels are in given ranges.)</li></ul><p>When composing  multiple sentences, use `|` as the delimiter, see given examples below.</p><p><b>Note</b>: When using more steps, the results can improve.</p>'
+
+iface = gr.Interface(compose,
+                     inputs=[
+                         "text",
+                         gr.Radio(['Stable_Diffusion_1v_4', 'GLIDE', 'CLEVR Objects'], type="value", label='version'),
+                         gr.Slider(2, 15),
+                         gr.Slider(10, 200)
+                     ],
+                     outputs='image',
                      title=title, description=description, examples=examples)
 
 iface.launch()