Update

app.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import functools
 import os
 import random
 import shlex
@@ -65,26 +64,25 @@ def load_model(device: torch.device) -> nn.Module:
     return model
 
 
-def generate_z(z_dim: int, seed: int, device: torch.device) -> torch.Tensor:
-    return torch.from_numpy(np.random.RandomState(seed).randn(1, z_dim)).to(device).float()
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = load_model(device)
+
+
+def generate_z(z_dim: int, seed: int) -> torch.Tensor:
+    return torch.from_numpy(np.random.RandomState(seed).randn(1, z_dim)).float()
 
 
 @torch.inference_mode()
-def generate_image(
-    seed: int, truncation_psi: float, randomize_noise: bool, model: nn.Module, device: torch.device
-) -> np.ndarray:
+def generate_image(seed: int, truncation_psi: float, randomize_noise: bool) -> np.ndarray:
     seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
 
-    z = generate_z(model.style_dim, seed, device)
+    z = generate_z(model.style_dim, seed)
+    z = z.to(device)
     out, _ = model([z], truncation=truncation_psi, truncation_latent=model.latent_avg, randomize_noise=randomize_noise)
     out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
     return out[0].cpu().numpy()
 
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-model = load_model(device)
-fn = functools.partial(generate_image, model=model, device=device)
-
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     with gr.Row():
@@ -93,7 +91,7 @@ with gr.Blocks(css="style.css") as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             psi = gr.Slider(label="Truncation psi", minimum=0, maximum=2, step=0.05, value=0.7)
             randomize_noise = gr.Checkbox(label="Randomize Noise", value=False)
-            run_button = gr.Button("Run")
+            run_button = gr.Button()
         with gr.Column():
             result = gr.Image(label="Output")
     gr.Markdown(ARTICLE)
@@ -105,9 +103,11 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     ).then(
-        fn=fn,
+        fn=generate_image,
         inputs=[seed, psi, randomize_noise],
         outputs=result,
         api_name="run",
     )
-demo.queue(max_size=10).launch()
+
+if __name__ == "__main__":
+    demo.queue(max_size=10).launch()