Update

README.md

@@ -4,7 +4,7 @@ emoji: 🏢
 colorFrom: blue
 colorTo: gray
 sdk: gradio
-sdk_version: 3.36.1
+sdk_version: 4.36.0
 app_file: app.py
 pinned: false
 suggested_hardware: t4-small

app.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import functools
 import pickle
 import sys
 
@@ -31,19 +30,25 @@ def load_model(file_name: str, 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("stylegan_human_v2_1024.pkl", 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_interpolated_images(
-    seed0: int, psi0: float, seed1: int, psi1: float, num_intermediate: int, model: nn.Module, device: torch.device
+    seed0: int, psi0: float, seed1: int, psi1: float, num_intermediate: int
 ) -> list[np.ndarray]:
     seed0 = int(np.clip(seed0, 0, np.iinfo(np.uint32).max))
     seed1 = int(np.clip(seed1, 0, np.iinfo(np.uint32).max))
 
-    z0 = generate_z(model.z_dim, seed0, device)
-    z1 = generate_z(model.z_dim, seed1, device)
+    z0 = generate_z(model.z_dim, seed0)
+    z1 = generate_z(model.z_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)]
@@ -61,12 +66,8 @@ def generate_interpolated_images(
     return res
 
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-model = load_model("stylegan_human_v2_1024.pkl", device)
-fn = functools.partial(generate_interpolated_images, model=model, device=device)
-
-gr.Interface(
-    fn=fn,
+demo = gr.Interface(
+    fn=generate_interpolated_images,
     inputs=[
         gr.Slider(label="Seed 1", minimum=0, maximum=100000, step=1, value=0, randomize=True),
         gr.Slider(label="Truncation psi 1", minimum=0, maximum=2, step=0.05, value=0.7),
@@ -74,7 +75,10 @@ gr.Interface(
         gr.Slider(label="Truncation psi 2", minimum=0, maximum=2, step=0.05, value=0.7),
         gr.Slider(label="Number of Intermediate Frames", minimum=0, maximum=21, step=1, value=7),
     ],
-    outputs=gr.Gallery(label="Output Images", type="numpy"),
+    outputs=gr.Gallery(label="Output Images"),
     title=TITLE,
     description=DESCRIPTION,
-).queue(max_size=10).launch()
+)
+
+if __name__ == "__main__":
+    demo.queue(max_size=10).launch()

requirements.txt

@@ -1,5 +1,5 @@
-numpy==1.23.5
-Pillow==10.0.0
-scipy==1.10.1
+numpy==1.26.4
+Pillow==10.3.0
+scipy==1.13.1
 torch==2.0.1
 torchvision==0.15.2