Update app.py

app.py

@@ -13,35 +13,47 @@ from PIL import Image
 import warnings
 warnings.filterwarnings('ignore')
 
-path_to_model = hf_hub_download(repo_id="opetrova/face-frontalization", filename="generator_v0.pt")
+# モデルのダウンロード
+path_to_model = hf_hub_download(
+    repo_id="opetrova/face-frontalization",
+    filename="generator_v0.pt"
+)
 
-# Download network.py into the current directory
+# network.py をカレントディレクトリにダウンロード
 network_url = hf_hub_url(repo_id="opetrova/face-frontalization", filename="network.py")
 r = requests.get(network_url, allow_redirects=True)
 open('network.py', 'wb').write(r.content)
 
-saved_model = torch.load(path_to_model, map_location=torch.device('cpu'))
+# PyTorch 2.6 以降は weights_only=False を指定しないとエラーになる
+saved_model = torch.load(path_to_model, map_location=torch.device("cpu"), weights_only=False)
 
 def frontalize(image):
-    
-    # Convert the test image to a [1, 3, 128, 128]-shaped torch tensor
-    # (as required by the frontalization model)
-    preprocess = transforms.Compose((transforms.ToPILImage(), 
-                                     transforms.Resize(size = (128, 128)), 
-                                     transforms.ToTensor()))
+    # 画像を [1, 3, 128, 128] tensor に変換
+    preprocess = transforms.Compose((
+        transforms.ToPILImage(),
+        transforms.Resize(size=(128, 128)),
+        transforms.ToTensor(),
+    ))
     input_tensor = torch.unsqueeze(preprocess(image), 0)
-    
-    # Use the saved model to generate an output (whose values go between -1 and 1, 
-    # and this will need to get fixed before the output is displayed)
-    generated_image = saved_model(Variable(input_tensor.type('torch.FloatTensor')))
+
+    # 推論
+    generated_image = saved_model(Variable(input_tensor.type(torch.FloatTensor)))
     generated_image = generated_image.detach().squeeze().permute(1, 2, 0).numpy()
-    generated_image = (generated_image + 1.0) / 2.0
-    
+    generated_image = (generated_image + 1.0) / 2.0  # [-1,1] → [0,1]
+
     return generated_image
 
-iface = gr.Interface(frontalize, gr.inputs.Image(type="numpy"), "image",
-                                 title='Face Frontalization', 
-                                 description='PyTorch implementation of a supervised GAN (see <a href="https://blog.scaleway.com/gpu-instances-using-deep-learning-to-obtain-frontal-rendering-of-facial-images/">blog post</a>)',
-                                 examples=["amos.png", "clarissa.png"],
+# Gradio インターフェース
+iface = gr.Interface(
+    fn=frontalize,
+    inputs=gr.Image(type="numpy"),
+    outputs="image",
+    title="Face Frontalization",
+    description=(
+        'PyTorch implementation of a supervised GAN '
+        '(see <a href="https://blog.scaleway.com/gpu-instances-using-deep-learning-to-obtain-frontal-rendering-of-facial-images/">blog post</a>)'
+    ),
+    examples=["amos.png", "clarissa.png"],
 )
-iface.launch()                      
+
+iface.launch()