Actually perform predictions

app.py

@@ -1,14 +1,24 @@
 #!/usr/bin/env python
 # coding: utf-8
-import random
-from typing import List
 import gradio as gr
+import random
+import torch
 from collections import defaultdict
+from diffusers import DiffusionPipeline
 from functools import partial
+from itertools import zip_longest
+from typing import List
 from PIL import Image
 
 SELECT_LABEL = "Select as seed"
 
+MODEL_ID = "CompVis/ldm-text2im-large-256"
+STEPS = 50
+ETA = 0.3
+GUIDANCE_SCALE = 12
+
+ldm = DiffusionPipeline.from_pretrained(MODEL_ID)
+
 with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
     state = gr.Variable({
         'selected': -1,
@@ -17,15 +27,21 @@ with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
 
     def infer_seeded_image(prompt, seed):
         print(f"Prompt: {prompt}, seed: {seed}")
-        return Image.open(f"sample_outputs/seeded_1.png")
-
-    def infer_grid(prompt):
-        response = defaultdict(list)
-        for i in range(1, 7):
-            response["images"].append(Image.open(f"sample_outputs/{i}.png"))
-            response["seeds"].append(random.randint(0, 2 ** 32 -1))
-
-        return response["images"], response["seeds"]
+        images, _ = infer_grid(prompt, n=1, seeds=[seed])
+        return images[0]
+
+    def infer_grid(prompt, n=6, seeds=[]):
+        # Unfortunately we have to iterate instead requesting all images at once,
+        # because we have no way to get the generation seeds.
+        result = defaultdict(list)
+        for _, seed in zip_longest(range(n), seeds, fillvalue=None):
+            seed = random.randint(0, 2**32 - 1) if seed is None else seed
+            print(f"Setting seed {seed}")
+            _ = torch.manual_seed(seed)
+            images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6)["sample"]
+            result["images"].append(images[0])
+            result["seeds"].append(seed)
+        return result["images"], result["seeds"]
 
     def infer(prompt, state):
         """
@@ -51,15 +67,6 @@ with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
         boxes = [gr.Box.update(visible=v) for v in visible]
         return grid_images + [image_with_seed] + boxes + [state]
 
-
-    def image_block():
-        return gr.Image(
-            interactive=False, show_label=False
-        ).style(
-            # border = (True, True, False, True),
-            rounded = (True, True, False, False),
-        )
-
     def update_state(selected_index: int, value, state):
         if value == '':
             others_value = None
@@ -74,6 +81,14 @@ with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
         state["selected"] = -1
         return [''] * 6 + [gr.Box.update(visible=True), gr.Box.update(visible=False)] + [state]
 
+    def image_block():
+        return gr.Image(
+            interactive=False, show_label=False
+        ).style(
+            # border = (True, True, False, True),
+            rounded = (True, True, False, False),
+        )
+
     def radio_block():
         radio = gr.Radio(
             choices=[SELECT_LABEL], interactive=True, show_label=False,