update

app.py

@@ -13,6 +13,8 @@ import torch
 from time import time
 
 
+device = "cuda"
+
 def clear_image():
     return None
 
@@ -44,7 +46,6 @@ def sample_nearby_vectors(base_vector, epsilons=[0.3, 0.5, 0.7], percentages=[0.
 
 
 def initialize_models():
-    device = torch.device('cpu')
     pose_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/6DRepNet_300W_LP_AFLW2000.pth", local_dir="./")
     id_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/arcface-r100-glint360k.pth", local_dir="./")
     quality_model_weights = hf_hub_download(repo_id="BooBooWu/Vec2Face", filename="weights/magface-r100-glint360k.pth", local_dir="./")
@@ -87,7 +88,7 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension, pro
         feature = np.random.normal(0, 1.0, (1, 512))
     else:
         input_image = np.transpose(input_image, (2, 0, 1))
-        input_image = torch.from_numpy(input_image).unsqueeze(0).float()
+        input_image = torch.from_numpy(input_image).unsqueeze(0).float().to(device)
         input_image.div_(255).sub_(0.5).div_(0.5)
         feature = id_model(input_image).clone().detach().cpu().numpy()
 
@@ -99,14 +100,14 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension, pro
             updated_feature[0][dimension] = feature[0][dimension] + i
             updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
             features.append(updated_feature)
-        features = torch.tensor(np.vstack(features)).float()
+        features = torch.tensor(np.vstack(features)).float().to(device)
         if quality > 25:
             images, _ = generator.gen_image(features, quality_model, id_model, q_target=quality)
         else:
             _, _, images, *_ = generator(features)
     else:
         features = torch.repeat_interleave(torch.tensor(feature), 3, dim=0)
-        features = sample_nearby_vectors(features, [0.7], [1]).float()
+        features = sample_nearby_vectors(features, [0.7], [1]).float().to(device)
         if quality > 25 or pose > 20:
             images, _ = generator.gen_image(features, quality_model, id_model, pose_model=pose_model,
                                             q_target=quality, pose=pose, class_rep=features)
@@ -154,6 +155,7 @@ def toggle_inputs(use_pose):
     ]
 
 
+# 4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
 def main():
     with gr.Blocks() as demo:
         title = r"""
@@ -167,8 +169,7 @@ def main():
             1. Upload an image with a cropped face image or directly click <b>Submit</b> button, three images will be shown on the right. 
             2. You can control the image quality, image pose, and modify the values in the target dimensions to change the output images. 
             3. The output results will shown three results of dimension modification or pose images.
-            4. Since the demo is CPU-based, higher quality and larger pose need longer time to run.
-            5. Enjoy! 😊
+            4. Enjoy! 😊
             """
 
         gr.Markdown(title)