Update app.py

added radio i think

app.py

@@ -86,10 +86,9 @@ def generate_annotation(img, overlap=False, hand_encoding=False):
     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")
-model_type="Standard"
-if model_type=="Standard":
-    args = Namespace(
+
+
+std_args = Namespace(
         pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
         revision="non-ema",
         from_pt=True,
@@ -97,8 +96,7 @@ if model_type=="Standard":
         controlnet_revision=None,
         controlnet_from_pt=False,
     )
-if model_type=="Hand Encoding":
-    args = Namespace(
+enc_args = Namespace(
         pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
         revision="non-ema",
         from_pt=True,
@@ -107,35 +105,58 @@ if model_type=="Hand Encoding":
         controlnet_from_pt=False,
     )
 
-controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
-    args.controlnet_model_name_or_path,
-    revision=args.controlnet_revision,
-    from_pt=args.controlnet_from_pt,
+std_controlnet, std_controlnet_params = FlaxControlNetModel.from_pretrained(
+    std_args.controlnet_model_name_or_path,
+    revision=std_args.controlnet_revision,
+    from_pt=std_args.controlnet_from_pt,
+    dtype=jnp.float32, # jnp.bfloat16
+)
+enc_controlnet, enc_controlnet_params = FlaxControlNetModel.from_pretrained(
+    enc_args.controlnet_model_name_or_path,
+    revision=enc_args.controlnet_revision,
+    from_pt=enc_args.controlnet_from_pt,
     dtype=jnp.float32, # jnp.bfloat16
 )
 
-pipeline, pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
-    args.pretrained_model_name_or_path,
+
+
+std_pipeline, std_pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    std_args.pretrained_model_name_or_path,
     # tokenizer=tokenizer,
-    controlnet=controlnet,
+    controlnet=std_controlnet,
     safety_checker=None,
     dtype=jnp.float32, # jnp.bfloat16
-    revision=args.revision,
-    from_pt=args.from_pt,
+    revision=std_args.revision,
+    from_pt=std_args.from_pt,
 )
+enc_pipeline, enc_pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    enc_args.pretrained_model_name_or_path,
+    # tokenizer=tokenizer,
+    controlnet=enc_controlnet,
+    safety_checker=None,
+    dtype=jnp.float32, # jnp.bfloat16
+    revision=enc_args.revision,
+    from_pt=enc_args.from_pt,
+)
+
 
+std_pipeline_params["controlnet"] = std_controlnet_params
+std_pipeline_params = jax_utils.replicate(std_pipeline_params)
 
-pipeline_params["controlnet"] = controlnet_params
-pipeline_params = jax_utils.replicate(pipeline_params)
+enc_pipeline_params["controlnet"] = enc_controlnet_params
+enc_pipeline_params = jax_utils.replicate(enc_pipeline_params)
 
 rng = jax.random.PRNGKey(0)
 num_samples = jax.device_count()
 prng_seed = jax.random.split(rng, jax.device_count())
 
 
-def infer(prompt, negative_prompt, image):
+def infer(prompt, negative_prompt, image, model_type="Standard"):
     prompts = num_samples * [prompt]
-    prompt_ids = pipeline.prepare_text_inputs(prompts)
+    if model_type=="Standard":
+        prompt_ids = std_pipeline.prepare_text_inputs(prompts)
+    if model_type=="Hand Encoding":
+        prompt_ids = enc_pipeline.prepare_text_inputs(prompts)
     prompt_ids = shard(prompt_ids)
 
     if model_type=="Standard":
@@ -145,21 +166,39 @@ def infer(prompt, negative_prompt, image):
         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)
 
-    negative_prompt_ids = pipeline.prepare_text_inputs([negative_prompt] * num_samples)
-    negative_prompt_ids = shard(negative_prompt_ids)
+    if model_type=="Standard":
+        processed_image = std_pipeline.prepare_image_inputs(num_samples * [validation_image])
+        processed_image = shard(processed_image)
+
+        negative_prompt_ids = std_pipeline.prepare_text_inputs([negative_prompt] * num_samples)
+        negative_prompt_ids = shard(negative_prompt_ids)
+
+        images = std_pipeline(
+            prompt_ids=prompt_ids,
+            image=processed_image,
+            params=std_pipeline_params,
+            prng_seed=prng_seed,
+            num_inference_steps=50,
+            neg_prompt_ids=negative_prompt_ids,
+            jit=True,
+        ).images
+    if model_type=="Hand Encoding":
+        processed_image = enc_pipeline.prepare_image_inputs(num_samples * [validation_image])
+        processed_image = shard(processed_image)
 
-    images = pipeline(
-        prompt_ids=prompt_ids,
-        image=processed_image,
-        params=pipeline_params,
-        prng_seed=prng_seed,
-        num_inference_steps=50,
-        neg_prompt_ids=negative_prompt_ids,
-        jit=True,
-    ).images
+        negative_prompt_ids = enc_pipeline.prepare_text_inputs([negative_prompt] * num_samples)
+        negative_prompt_ids = shard(negative_prompt_ids)
+
+        images = enc_pipeline(
+            prompt_ids=prompt_ids,
+            image=processed_image,
+            params=enc_pipeline_params,
+            prng_seed=prng_seed,
+            num_inference_steps=50,
+            neg_prompt_ids=negative_prompt_ids,
+            jit=True,
+        ).images
 
 
     images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
@@ -176,16 +215,15 @@ with gr.Blocks(theme='gradio/soft') as demo:
     Model1 can be found at [https://huggingface.co/Vincent-luo/controlnet-hands](https://huggingface.co/Vincent-luo/controlnet-hands)
     
     Model2 can be found at [https://huggingface.co/MakiPan/controlnet-encoded-hands-130k/ ](https://huggingface.co/MakiPan/controlnet-encoded-hands-130k/)
-
     Dataset1 can be found at [https://huggingface.co/datasets/MakiPan/hagrid250k-blip2](https://huggingface.co/datasets/MakiPan/hagrid250k-blip2)
     
     Dataset2 can be found at [https://huggingface.co/datasets/MakiPan/hagrid-hand-enc-250k](https://huggingface.co/datasets/MakiPan/hagrid-hand-enc-250k)
     
     Preprocessing1 can be found at [https://github.com/Maki-DS/Jax-Controlnet-hand-training/blob/main/normal-preprocessing.py](https://github.com/Maki-DS/Jax-Controlnet-hand-training/blob/main/normal-preprocessing.py)
-
     Preprocessing2 can be found at [https://github.com/Maki-DS/Jax-Controlnet-hand-training/blob/main/Hand-encoded-preprocessing.py](https://github.com/Maki-DS/Jax-Controlnet-hand-training/blob/main/Hand-encoded-preprocessing.py)
     """)
-    
+    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")
+
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(label="Prompt")
@@ -227,13 +265,13 @@ with gr.Blocks(theme='gradio/soft') as demo:
                "example4.png"
             ],
         ],
-        inputs=[prompt_input, negative_prompt, input_image],
+        inputs=[prompt_input, negative_prompt, input_image, model_type],
         outputs=[output_image],
         fn=infer,
         cache_examples=True,
     )
 
-    inputs = [prompt_input, negative_prompt, input_image]
+    inputs = [prompt_input, negative_prompt, input_image, model_type]
     submit_btn.click(fn=infer, inputs=inputs, outputs=[output_image])
 
 demo.launch()