Update app.py

app.py

@@ -69,14 +69,15 @@ def on_change_event(app_state):
         step = app_state['step']
         label = f'Reconstructed image from the latent state at step {step}. It will get better :)'
         print(f'Updating the image:! {app_state}')
-        return gr.update(value=img, label=label)
+        return gr.update(value=img, label=label), gr.update(value=app_state['img_list'], label='intermediate steps')
     else:
-        return gr.update(label='Illustration will appear here soon')
+        return gr.update(label='Illustration will appear here soon'), gr.update(label='images list')
 
 with gr.Blocks() as demo:
     
     def generate_image(prompt, inference_steps, app_state):
         app_state['running'] = True
+        app_state['img_list'] = []
         def callback(step, ts, latents):
             print (f'In Callback on {step} {ts} !')
             latents = 1 / 0.18215 * latents
@@ -86,17 +87,20 @@ with gr.Blocks() as demo:
             res = pipe.numpy_to_pil(res)[0]
             app_state['img'] = res
             app_state['step'] = step
+            app_state['img_list'].append(res)
             print (f'In Callback on {app_state} Done!')
         
         prompt = prompt + ' masterpiece charcoal pencil art lord of the rings illustration'
-        img = pipe(prompt, height=512, width=512, num_inference_steps=inference_steps, callback=callback, callback_steps=3)
+        img = pipe(prompt, height=512, width=512, num_inference_steps=inference_steps, callback=callback, callback_steps=1)
         app_state['running'] = False
         app_state['img'] = None
         return gr.update(value=img.images[0], label='Generated image')
     
     app_state = gr.State({'img': None, 
                           'step':0, 
-                          'running':False})
+                          'running':False,
+                           'img_list': []
+                         })
     title = gr.Markdown('## Lord of the rings app')
     description = gr.Markdown(f'#### A Lord of the rings inspired app that combines text and image generation.'
                               f' The language modeling is done by fine tuning distilgpt2 on the LOTR trilogy.'
@@ -112,23 +116,26 @@ with gr.Blocks() as demo:
                                   ' but here you can see here what is generated from the latent state of the diffuser every few steps.'
                                   ' Usually there is a significant improvement around step 12 that yields much better result')
     image = gr.Image(label='Illustration for your story', show_label=True)
+    gallery = gr.Gallery()
+    gallery.style(grid=[4])
     
     inference_steps = gr.Slider(5, 30, 
                                 value=20, 
                                 step=1, 
                                 visible=True,
-                                label=f"Num inference steps (more steps makes a better image but takes more time)")
+                                label=f"Num inference steps (more steps yields a better image but takes more time)")
     
     
     bt_make_text.click(fn=generate_story, inputs=prompt, outputs=[story, summary, bt_make_image])
     bt_make_image.click(fn=generate_image, inputs=[summary, inference_steps, app_state], outputs=image)
     
     eventslider = gr.Slider(visible=False)
-    dep = demo.load(on_change_event, app_state, image, every=10)
-    eventslider.change(fn=on_change_event, inputs=[app_state], outputs=[image], every=10, cancels=[dep])
+    dep = demo.load(on_change_event, app_state, [image, gallery], every=10)
+    eventslider.change(fn=on_change_event, inputs=[app_state], outputs=[image, gallery], every=10, cancels=[dep])
     
 
 if READ_TOKEN:
     demo.queue().launch()
 else:
     demo.queue().launch(share=True, debug=True)
+