last/first frame output, reordering some options, examples

.gitattributes

@@ -1,5 +1,6 @@
 *.webp filter=lfs diff=lfs merge=lfs -text
 *.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,5 +1,6 @@
 
 import os
+import json
 from io import BytesIO
 import base64
 from functools import partial
@@ -16,8 +17,6 @@ from makeavid_sd.inference import (
 print(os.environ.get('XLA_PYTHON_CLIENT_PREALLOCATE', 'NotSet'))
 print(os.environ.get('XLA_PYTHON_CLIENT_ALLOCATOR', 'NotSet'))
 
-_preheat: bool = False
-
 _seen_compilations = set()
 
 _model = InferenceUNetPseudo3D(
@@ -30,9 +29,20 @@ if _model.failed != False:
     trace = f'```{_model.failed}```'
     with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled = False) as demo:
         exception = gr.Markdown(trace)
-
     demo.launch()
 
+_examples = []
+_expath = 'examples'
+for x in os.listdir(_expath):
+    with open(os.path.join(_expath, x, 'params.json'), 'r') as f:
+        ex = json.load(f)
+    ex['image_input'] = None
+    if os.path.isfile(os.path.join(_expath, x, 'input.png')):
+        ex['image_input'] = os.path.join(_expath, x, 'input.png')
+    ex['image_output'] = os.path.join(_expath, x, 'output.gif')
+    _examples.append(ex)
+
+
 _output_formats = (
     'webp', 'gif'
 )
@@ -100,10 +110,17 @@ def generate(
             duration = round(1000 / fps),
             allow_mixed = True
     )
-    data = base64.b64encode(buffer.getvalue()).decode()
+    data = f'data:image/{output_format};base64,' + base64.b64encode(buffer.getvalue()).decode()
+    buffer.close()
+    buffer = BytesIO()
+    images[-1].save(buffer, format ='png')
+    last_data = f'data:image/png;base64,' + base64.b64encode(buffer.getvalue()).decode()
+    buffer.close()
+    buffer = BytesIO()
+    images[0].save(buffer, format ='png')
+    first_data = f'data:image/png;base64,' + base64.b64encode(buffer.getvalue()).decode()
     buffer.close()
-    data = f'data:image/{output_format};base64,' + data
-    return data
+    return data, last_data, first_data
 
 def check_if_compiled(hint_image, inference_steps, height, width, num_frames, scheduler_type, message):
     height = int(height)
@@ -116,27 +133,6 @@ def check_if_compiled(hint_image, inference_steps, height, width, num_frames, sc
     else:
         return  f"""{message}"""
 
-if _preheat:
-    print('\npreheating the oven')
-    generate(
-            prompt = 'preheating the oven',
-            neg_prompt = '',
-            image = None,
-            inference_steps = 20,
-            cfg = 12.0,
-            seed = 0
-    )
-    print('Entertaining the guests with sailor songs played on an old piano.')
-    dada = generate(
-            prompt = 'Entertaining the guests with sailor songs played on an old harmonium.',
-            neg_prompt = '',
-            image = Image.new('RGB', size = (512, 512), color = (0, 0, 0)),
-            inference_steps = 20,
-            cfg = 12.0,
-            seed = 0
-    )
-    print('dinner is ready\n')
-
 with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled = False) as demo:
     variant = 'panel'
     with gr.Row():
@@ -283,6 +279,52 @@ with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled =
                     value = 'example.gif',
                     interactive = False
             )
+            tips = gr.Markdown('🤫 *Secret tip*: take the last frame as input for the next generation.')
+            with gr.Row():
+                last_frame_output = gr.Image(
+                        label = 'Last frame',
+                        interactive = False
+                )
+                first_frame_output = gr.Image(
+                        label = 'First frame',
+                        interactive = False
+                )
+    examples_lst = []
+    for x in _examples:
+        examples_lst.append([
+            x['image_output'],
+            x['prompt'],
+            x['neg_prompt'],
+            x['image_input'],
+            x['cfg'],
+            x['cfg_image'],
+            x['seed'],
+            x['fps'],
+            x['num_frames'],
+            x['height'],
+            x['width'],
+            x['scheduler'],
+            x['format']
+        ])
+    examples = gr.Examples(
+            examples = examples_lst,
+            inputs = [
+                    image_output,
+                    prompt_input,
+                    neg_prompt_input,
+                    image_input,
+                    cfg_input,
+                    cfg_image_input,
+                    seed_input,
+                    fps_input,
+                    num_frames_input,
+                    height_input,
+                    width_input,
+                    scheduler_input,
+                    output_format
+            ],
+            postprocess = False
+    )
     #trigger_inputs = [ image_input, inference_steps_input, height_input, width_input, num_frames_input, scheduler_input ]
     #trigger_check_fun = partial(check_if_compiled, message = 'Current parameters need compilation.')
     #height_input.change(fn = trigger_check_fun, inputs = trigger_inputs, outputs = will_trigger)
@@ -308,14 +350,11 @@ with gr.Blocks(title = 'Make-A-Video Stable Diffusion JAX', analytics_enabled =
                     scheduler_input,
                     output_format
             ],
-            outputs = image_output,
+            outputs = [ image_output, last_frame_output, first_frame_output ],
             postprocess = False
     )
     #cancel_button.click(fn = lambda: None, cancels = ev)
 
-demo.queue(concurrency_count = 1, max_size = 12)
+demo.queue(concurrency_count = 1, max_size = 10)
 demo.launch()
 
-# Photorealistic fantasy oil painting of  the angry minotaur in a threatening pose by Randy Vargas.
-# A girl is dancing by a beautiful lake by sophie anderson and greg rutkowski and alphonse mucha.
-# They are dancing in the club but everybody is a 3d cg  hairy monster wearing a hairy costume.