feat: 4 tasks

app.py

@@ -4,6 +4,9 @@ from PIL import Image
 from inference import generate_image
 
 
+TASK_TO_INDEX = {"Task 1": 0, "Task 2": 1, "Task 3": 2, "Task 4": 3}
+
+
 def update_reference_image(choice: int) -> tuple[str, int, str]:
     """
     Update the reference image display based on radio button selection
@@ -24,7 +27,15 @@ def process_coord_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
     return generate_image(image_idx, x, y)
 
 
-with gr.Blocks() as demo:
+with gr.Blocks(
+    css="""
+    .radio-container { 
+        width: 450px !important;
+        margin-left: auto !important;
+        margin-right: auto !important;
+    }
+"""
+) as demo:
     gr.Markdown(
         """
     # Interactive Image Generation
@@ -38,19 +49,20 @@ with gr.Blocks() as demo:
             # State variable to track selected image index
             selected_idx = gr.State(value=0)
 
-            # Radio buttons for task selection
-            task_select = gr.Radio(
-                choices=["Task 1", "Task 2"],
-                value="Task 1",
-                label="Select Task",
-                interactive=True,
-            )
+            # Radio buttons with container class
+            with gr.Column(elem_classes="radio-container"):
+                task_select = gr.Radio(
+                    choices=["Task 1", "Task 2", "Task 3", "Task 4"],
+                    value="Task 1",
+                    label="Select Task",
+                    interactive=True,
+                )
 
             # Single reference image component that updates based on selection
             reference_image = gr.Image(
                 value="imgs/pattern_0.png",
                 show_label=False,
-                interactive=False,  # No need for click interaction now
+                interactive=False,
                 height=300,
                 width=450,
             )
@@ -60,7 +72,7 @@ with gr.Blocks() as demo:
             # Coordinate selector with dynamic background
             coord_selector = gr.Image(
                 value="imgs/heatmap_0.png",
-                label="Click to select (x, y) coordinates in the latent space",
+                label="Click to select (x, y) coordinates",
                 show_label=True,
                 interactive=True,
                 height=400,
@@ -72,7 +84,7 @@ with gr.Blocks() as demo:
 
     # Handle radio button selection
     task_select.change(
-        fn=lambda x: update_reference_image(0 if x == "Task 1" else 1),
+        fn=lambda x: update_reference_image(TASK_TO_INDEX[x]),
         inputs=[task_select],
         outputs=[reference_image, selected_idx, coord_selector],
     )

inference.py

@@ -23,7 +23,7 @@ from utils import patch_target, ax_to_pil
 
 
 checkpoint_name = "quiet-thunder-789--checkpoint:v0"
-BLUE_LOCATION_INPUTS = {0: 13, 1: 9}
+BLUE_LOCATION_INPUTS = {0: 13, 1: 9, 2: 9, 3: 6}
 
 local_dir = snapshot_download(repo_id="clement-bonnet/lpn-2d", allow_patterns=f"{checkpoint_name}/*")
 with open(f"{local_dir}/{checkpoint_name}/config.yaml", "r") as f: