Update app.py

app.py

@@ -1,5 +1,6 @@
 import time
 import os
+import PIL
 import gradio as gr
 
 import torch
@@ -19,7 +20,7 @@ if has_cuda:
     pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=READ_TOKEN)
     device = "cuda"
 else:
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", use_auth_token=READ_TOKEN)
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=READ_TOKEN)
     device = "cpu"
     
 pipe.to(device)
@@ -33,6 +34,8 @@ tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
 
 summarizer = pipeline("summarization")
 
+#######################################################
+
 def break_until_dot(txt):
     return txt.rsplit('.', 1)[0] + '.'
 
@@ -42,7 +45,8 @@ def generate(prompt):
 
     outputs = model.generate(
         input_ids=input_ids, 
-        max_length=180, 
+        max_length=120,
+        min_length=50,
         temperature=0.7, 
         num_return_sequences=3, 
         do_sample=True
@@ -57,22 +61,22 @@ def generate_story(prompt):
     summary = break_until_dot(summary)
     return story, summary, gr.update(visible=True)
 
-def on_change_event(app_state=None):
-    print(f'In change event!')
-    if app_state and app_state['running']:
+def on_change_event(app_state):
+    if app_state and app_state['running'] and app_state['img']:
         img = app_state['img']
         step = app_state['step']
         label = f'Reconstructed image from the latent state at step {step}'
+        print(f'Updating the image:! {app_state}')
         return gr.update(value=img, label=label)
     else:
-        return None
+        return gr.update()
 
 with gr.Blocks() as demo:
     
     def generate_image(prompt, inference_steps, app_state):
         app_state['running'] = True
         def callback(step, ts, latents):
-            print (f'In Callback on {step}!')
+            print (f'In Callback on {step} {ts} !')
             latents = 1 / 0.18215 * latents
             res = pipe.vae.decode(latents).sample
             res = (res / 2 + 0.5).clamp(0, 1)
@@ -80,9 +84,10 @@ with gr.Blocks() as demo:
             res = pipe.numpy_to_pil(res)[0]
             app_state['img'] = res
             app_state['step'] = step
+            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=5)
+        img = pipe(prompt, height=512, width=512, num_inference_steps=inference_steps, callback=callback, callback_steps=2)
         app_state['running'] = False
         return gr.update(value=img.images[0], label='Generated image')
     
@@ -103,7 +108,7 @@ with gr.Blocks() as demo:
     img_description = gr.Markdown('Image generation take some time'
                                   ' but here you can see the what is generated from the latent state of the diffuser every few steps.'
                                   ' Usually there is a significant improvement around step 15, that yields much better result')
-    image = gr.Image(label='Illustration for your story', shape=(512, 512), show_label=True)
+    image = gr.Image(label='Illustration for your story', show_label=True)
     
     inference_steps = gr.Slider(5, 30, 
                                 value=15, 
@@ -115,15 +120,12 @@ with gr.Blocks() as demo:
     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)
     
-#     bt_boo = gr.Button("Click me")
-#     bt_boo.click(fn=on_change_event, inputs=app_state, outputs=image, every=1)
-#     eventslider = gr.Slider(label='Boooo!')
-#     dep = demo.load(on_change_event, None, None, every=1)
-#     eventslider.change(fn=on_change_event, inputs=[app_state], outputs=[image], every=1, cancels=[dep])
-    inference_steps.change(fn=on_change_event, inputs=[app_state], outputs=[image], every=1)
+    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])
     
 
 if READ_TOKEN:
     demo.queue().launch()
 else:
-    demo.queue().launch(share=True, debug=True)
+    demo.queue().launch(share=True, debug=True)