Added support for witch and werewolf

app.py

@@ -16,8 +16,12 @@ from faceparsing2 import get_face_mask
 # model.load_state_dict(torch.load('models/your_pix2pixhd_model.pth'))
 # model.eval()
 
-model_path = hf_hub_download(repo_id="jbrownkramer/makemeazombie", filename="smaller512x512_32bit.onnx")
-ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])
+model_map = {"zombie":"smaller512x512_32bit.onnx","witch":"witch.onnx","werewolf":"werewolf.onnx"}
+
+inference_session_map = {}
+for model_name, model_filename in model_map.items():
+    model_path = hf_hub_download(repo_id="jbrownkramer/makemeazombie", filename=model_filename)
+    inference_session_map[model_name] = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])
 
 # --- 2. Define the prediction function ---
 # def predict(input_image):
@@ -34,17 +38,17 @@ ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider
 
 #     # return output_image
 
-def predict(input_image, mode):
+def predict(input_image, mode, model_name):
     input_image = input_image.convert("RGB")   
     if mode == "Classic":
         # Use the transition_onnx function for side-by-side comparison
-        zombie_image = zombie.transition_onnx(input_image, ort_session)
+        zombie_image = zombie.transition_onnx(input_image, inference_session_map[model_name])
         if zombie_image is None:
             return "No face found"
         return zombie_image
     elif mode == "In Place":
         im_array = np.array(input_image)
-        zombie_image = zombie.make_faces_zombie_from_array(im_array, None, ort_session)
+        zombie_image = zombie.make_faces_zombie_from_array(im_array, None, inference_session_map[model_name])
         if zombie_image is None:
             return "No face found"
         return zombie_image
@@ -61,7 +65,8 @@ demo = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil", label="Input Image"),
-        gr.Dropdown(choices=["Classic", "In Place"], value="Classic", label="Mode")
+        gr.Dropdown(choices=["Classic", "In Place"], value="Classic", label="Mode"),
+        gr.Dropdown(choices=["zombie", "witch", "werewolf"], value="zombie", label="Model")
     ],
     outputs=gr.Image(type="pil", label="Output Image"),
     title=title,

requirements.txt

@@ -1,5 +1,5 @@
 gradio
-onnxruntime
+onnxruntime-gpu
 opencv-python
 numpy
 mediapipe