Update

app.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import functools
 import os
 import random
 import shlex
@@ -58,12 +57,16 @@ 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(model: nn.Module, z: torch.Tensor, truncation_psi: float, randomize_noise: bool) -> np.ndarray:
+def generate_image(z: torch.Tensor, truncation_psi: float, randomize_noise: bool) -> np.ndarray:
     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()
@@ -77,14 +80,15 @@ def generate_interpolated_images(
     psi0: float,
     psi1: float,
     randomize_noise: bool,
-    model: nn.Module,
-    device: torch.device,
 ) -> list[np.ndarray]:
     seed0 = int(np.clip(seed0, 0, MAX_SEED))
     seed1 = int(np.clip(seed1, 0, MAX_SEED))
 
-    z0 = generate_z(model.style_dim, seed0, device)
-    z1 = generate_z(model.style_dim, seed1, device)
+    z0 = generate_z(model.style_dim, seed0)
+    z1 = generate_z(model.style_dim, seed1)
+    z0 = z0.to(device)
+    z1 = z1.to(device)
+
     vec = z1 - z0
     dvec = vec / (num_intermediate + 1)
     zs = [z0 + dvec * i for i in range(num_intermediate + 2)]
@@ -92,15 +96,11 @@ def generate_interpolated_images(
     psis = [psi0 + dpsi * i for i in range(num_intermediate + 2)]
     res = []
     for z, psi in zip(zs, psis):
-        out = generate_image(model, z, psi, randomize_noise)
+        out = generate_image(z, psi, randomize_noise)
         res.append(out)
     return res
 
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-model = load_model(device)
-fn = functools.partial(generate_interpolated_images, model=model, device=device)
-
 examples = [
     [29703, 55376, 3, 0.7, 0.7, False],
     [34141, 36864, 5, 0.7, 0.7, False],
@@ -141,13 +141,15 @@ with gr.Blocks(css="style.css") as demo:
         examples=examples,
         inputs=inputs,
         outputs=result,
-        fn=fn,
+        fn=generate_interpolated_images,
         cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
     )
     run_button.click(
-        fn=fn,
+        fn=generate_interpolated_images,
         inputs=inputs,
         outputs=result,
         api_name="run",
     )
-demo.queue(max_size=10).launch()
+
+if __name__ == "__main__":
+    demo.queue(max_size=10).launch()