Update app.py

added with gr.row for radio and added conditionals depending on model selection

app.py

@@ -85,17 +85,27 @@ def generate_annotation(img, overlap=False, hand_encoding=False):
     # STEP 5: Process the classification result. In this case, visualize it.
     annotated_image = draw_landmarks_on_image(image.numpy_view(), detection_result, overlap=overlap, hand_encoding=hand_encoding)
     return annotated_image
-
-model_type = gr.Radio(["Standard", "Hand Encoding"], label="Model preprocessing", info="We developed two models, one with standard mediapipe landmarks, and one with different (but similar) coloring on palm landmards to distinguish left and right")
-
-args = Namespace(
-    pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
-    revision="non-ema",
-    from_pt=True,
-    controlnet_model_name_or_path="Vincent-luo/controlnet-hands",
-    controlnet_revision=None,
-    controlnet_from_pt=False,
-)
+with gr.Row():
+    model_type = gr.Radio(["Standard", "Hand Encoding"], label="Model preprocessing", info="We developed two models, one with standard mediapipe landmarks, and one with different (but similar) coloring on palm landmards to distinguish left and right")
+
+if model_type=="Standard":
+    args = Namespace(
+        pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
+        revision="non-ema",
+        from_pt=True,
+        controlnet_model_name_or_path="Vincent-luo/controlnet-hands",
+        controlnet_revision=None,
+        controlnet_from_pt=False,
+    )
+if model_type=="Hand Encoding":
+    args = Namespace(
+        pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
+        revision="non-ema",
+        from_pt=True,
+        controlnet_model_name_or_path="MakiPan/controlnet-encoded-hands-130k",
+        controlnet_revision=None,
+        controlnet_from_pt=False,
+    )
 
 controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
     args.controlnet_model_name_or_path,
@@ -128,7 +138,12 @@ def infer(prompt, negative_prompt, image):
     prompt_ids = pipeline.prepare_text_inputs(prompts)
     prompt_ids = shard(prompt_ids)
 
-    annotated_image = generate_annotation(image)
+    if model_type=="Standard":
+        annotated_image = generate_annotation(image, overlap=False, hand_encoding=False)
+        overlap_image = generate_annotation(image, overlap=True, hand_encoding=False)
+    if model_type=="Hand Encoding":
+        annotated_image = generate_annotation(image, overlap=False, hand_encoding=True)
+        overlap_image = generate_annotation(image, overlap=True, hand_encoding=True)
     validation_image = Image.fromarray(annotated_image).convert("RGB")
     processed_image = pipeline.prepare_image_inputs(num_samples * [validation_image])
     processed_image = shard(processed_image)
@@ -150,7 +165,7 @@ def infer(prompt, negative_prompt, image):
     images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
 
     results = [i for i in images]
-    return [annotated_image] + results
+    return [annotated_image, overlap_image] + results
 
 
 with gr.Blocks(theme='gradio/soft') as demo: