Update app.py

app.py

@@ -82,33 +82,44 @@ for m in G.modules():
             m.recompute_scale_factor = None
 
 
-def generate_image(seed: int, truncation_psi: float):
+def generate_image(seed: int):
     label = torch.zeros([1, G.c_dim], device=device)
     z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device).float()
     if hasattr(G.synthesis, 'input'):
         m = make_transform((0, 0), 0)
         m = np.linalg.inv(m)
         G.synthesis.input.transform.copy_(torch.from_numpy(m))
-    img = G(z, label, truncation_psi=truncation_psi, noise_mode='const')
+    img = G(z, label, truncation_psi=1, noise_mode='const')
     img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
     return PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
 
 
-def inference(seed, truncation_psi):
+def inference(seed):
     seed = int(seed)
-    image = generate_image(seed, float(truncation_psi))
+    image = generate_image(seed)
     return image, str(seed)
 
 
+def run_and_log(seed, history):
+    image, seed_str = inference(seed)
+    history = (history or []) + [(image, f"seed={int(seed)}")]
+    return image, seed_str, history, history
+
+
 def pick_random_seed():
     return random.randint(0, 5000)
 
 
+def clear_history():
+    return [], []
+
+
 header_md = """
 # Projected GAN — Pokémon
 
 Generate Pokémon with the pretrained [Projected GAN](https://github.com/autonomousvision/projected_gan)
-from Sauer et al., NeurIPS 2021. Pick a seed for reproducibility, or roll the dice.
+from Sauer et al., NeurIPS 2021. Pick a seed, roll the dice, or click a curated example.
+Each Generate adds to your session history below.
 """
 
 footer_md = """
@@ -117,48 +128,62 @@ Paper: [Projected GANs Converge Faster](http://www.cvlibs.net/publications/Sauer
 Code: [autonomousvision/projected_gan](https://github.com/autonomousvision/projected_gan)
 """
 
-with gr.Blocks(theme=gr.themes.Soft(), title="Projected GAN — Pokémon") as demo:
+custom_css = """
+.gradio-container { max-width: 1100px !important; }
+#main-image img { object-fit: contain; }
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="Projected GAN — Pokémon") as demo:
     gr.Markdown(header_md)
 
-    with gr.Row():
-        with gr.Column(scale=1):
+    gallery_state = gr.State([])
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=2):
             seed_slider = gr.Slider(
                 label="Seed",
                 minimum=0, maximum=5000, step=1, value=42,
-                info="Same seed + same truncation = same image.",
-            )
-            trunc_slider = gr.Slider(
-                label="Truncation ψ",
-                minimum=0.0, maximum=1.0, step=0.05, value=1.0,
-                info="Lower = more typical / less varied. Higher = more diverse / weirder.",
+                info="Same seed = same image.",
             )
             with gr.Row():
                 random_btn = gr.Button("🎲 Random seed")
                 generate_btn = gr.Button("Generate", variant="primary")
-
-        with gr.Column(scale=1):
-            image_out = gr.Image(label="Output", type="pil")
             seed_used = gr.Textbox(label="Seed used", interactive=False)
 
+        with gr.Column(scale=3):
+            image_out = gr.Image(label="Output", type="pil", height=512, elem_id="main-image")
+
     gr.Examples(
-        examples=[
-            [0, 1.0], [1, 1.0], [10, 1.0], [20, 1.0], [30, 1.0],
-            [42, 1.0], [50, 1.0], [60, 1.0], [77, 1.0], [102, 1.0],
-        ],
-        inputs=[seed_slider, trunc_slider],
+        examples=[[0], [1], [10], [20], [30], [42], [50], [60], [77], [102]],
+        inputs=[seed_slider],
         outputs=[image_out, seed_used],
         fn=inference,
         cache_examples=False,
         label="Curated seeds",
     )
 
+    with gr.Row():
+        gr.Markdown("### Session history")
+    gallery = gr.Gallery(
+        label=None, show_label=False, columns=4, height=340,
+    )
+    with gr.Row():
+        clear_btn = gr.Button("🗑 Clear history", size="sm")
+
     gr.Markdown(footer_md)
 
     random_btn.click(fn=pick_random_seed, outputs=seed_slider)
     generate_btn.click(
-        fn=inference,
-        inputs=[seed_slider, trunc_slider],
-        outputs=[image_out, seed_used],
+        fn=run_and_log,
+        inputs=[seed_slider, gallery_state],
+        outputs=[image_out, seed_used, gallery_state, gallery],
+    )
+    clear_btn.click(fn=clear_history, outputs=[gallery_state, gallery])
+
+    demo.load(
+        fn=run_and_log,
+        inputs=[seed_slider, gallery_state],
+        outputs=[image_out, seed_used, gallery_state, gallery],
     )
 
 demo.queue().launch()