feat: code wip

app.py

@@ -1,122 +1,135 @@
-import os
+# import gradio as gr
+# import numpy as np
+# from PIL import Image
+
+# from inference import generate_image
+
+
+# # Create a square image for the coordinate selector
+# def create_selector_image():
+#     # Create a white square with black border
+#     size = 400
+#     border = 2
+#     img = np.ones((size, size, 3), dtype=np.uint8) * 255
+#     # Add black border
+#     img[:border, :] = 0  # top
+#     img[-border:, :] = 0  # bottom
+#     img[:, :border] = 0  # left
+#     img[:, -border:] = 0  # right
+#     return Image.fromarray(img)
+
+
+# def process_click(image_idx: int, x: int, y: int) -> tuple[Image.Image, str]:
+#     """
+#     Process the click event on the coordinate selector
+#     """
+#     try:
+#         # Normalize coordinates to [0, 1]
+#         x_norm, y_norm = x / 400, y / 400  # Divide by image size (400x400)
+
+#         # Debug message
+#         debug_msg = f"Processing: image_idx={image_idx}, coordinates=({x_norm:.3f}, {y_norm:.3f})"
+#         print(debug_msg)
+
+#         # Generate image using the model
+#         generated_img = generate_image(image_idx, x_norm, y_norm)
+#         return generated_img, debug_msg
+#     except Exception as e:
+#         error_msg = f"Error: {str(e)}"
+#         print(error_msg)
+#         return None, error_msg
+
+
+# with gr.Blocks() as demo:
+#     gr.Markdown(
+#         """
+#     # Interactive Image Generation
+#     Choose a reference image and click on the coordinate selector to generate a new image.
+#     """
+#     )
+
+#     with gr.Row():
+#         # Left column: Reference images and coordinate selector
+#         with gr.Column(scale=1):
+#             # Radio buttons for image selection
+#             image_idx = gr.Radio(
+#                 choices=list(range(2)), value=0, label="Select Reference Image", type="index"
+#             )
+
+#             # Display reference images
+#             gallery = gr.Gallery(
+#                 value=["imgs/pattern_1.png", "imgs/pattern_2.png"],
+#                 columns=2,
+#                 rows=1,
+#                 height=500,
+#                 label="Different Tasks",
+#             )
+
+#             # Coordinate selector
+#             coord_selector = gr.Image(
+#                 value=create_selector_image(),
+#                 label="Click to select (x, y) coordinates",
+#                 show_label=True,
+#                 interactive=True,
+#                 height=400,
+#                 width=400,
+#             )
+
+#         # Right column: Generated image and debug info
+#         with gr.Column(scale=1):
+#             output_image = gr.Image(label="Generated Image", height=400)
+#             debug_text = gr.Textbox(label="Debug Info", interactive=False)
+
+#     # Handle click events using click instead of select
+#     coord_selector.click(
+#         fn=process_click,
+#         inputs=[image_idx, coord_selector],  # coord_selector will provide x, y coordinates
+#         outputs=[output_image, debug_text],
+#     )
+
+# if __name__ == "__main__":
+#     print("Starting Gradio app...")
+#     demo.launch(debug=True)
+
+
 import gradio as gr
 import numpy as np
 from PIL import Image
 
-# Placeholder for your actual model
-def generate_image(image_idx: int, x: float, y: float) -> Image.Image:
-    """
-    Replace this with your actual model inference
-    """
-    # This is just a placeholder - replace with your model
-    # Creating a simple gradient image as example output
-    width, height = 256, 256
-    gradient = np.zeros((height, width, 3), dtype=np.uint8)
-    gradient[:, :, 0] = np.linspace(0, 255 * x, width)
-    gradient[:, :, 1] = np.linspace(0, 255 * y, height)[:, np.newaxis]
-    gradient[:, :, 2] = image_idx * 30  # vary blue channel based on selected image
-    return Image.fromarray(gradient)
-
-def process_click(image_idx: int, evt: gr.SelectData) -> Image.Image:
-    """
-    Process the click event on the coordinate selector
-    """
-    # Extract coordinates from click event
-    x, y = evt.index[0], evt.index[1]
-    # Normalize coordinates to [0, 1]
-    x, y = x/100, y/100
-    # Generate image using the model
-    return generate_image(image_idx, x, y)
 
+def create_white_square(size=400):
+    # Create a white square image
+    print("Creating white square")
+    return np.full((size, size, 3), 255, dtype=np.uint8)
+
+
+def get_click_coordinates(evt: gr.SelectData):
+    # Get click coordinates
+    x, y = evt.index
+    print(f"Clicked at coordinates: x={x}, y={y}")
+    return f"Clicked at coordinates: x={x}, y={y}"
+
+
+# Create the interface
 with gr.Blocks() as demo:
-    gr.Markdown("""
-    # Interactive Image Generation
-    Choose a reference image and click on the coordinate selector to generate a new image.
-    """)
-    
-    with gr.Row():
-        # Left column: Reference images and coordinate selector
-        with gr.Column(scale=1):
-            # Radio buttons for image selection
-            image_idx = gr.Radio(
-                choices=[i for i in range(4)],  # Replace with your actual number of images
-                value=0,
-                label="Select Reference Image",
-                type="index"
-            )
-            
-            # Display reference images
-            gallery = gr.Gallery(
-                value=[
-                    "image_0.jpg",
-                    "image_0.jpg",
-                    "image_0.jpg",
-                    "image_0.jpg",
-                ],
-                columns=2,
-                rows=2,
-                height=300,
-                label="Reference Images"
-            )
-            
-            # Coordinate selector (displayed as heatmap for click interaction)
-            coord_selector = gr.Plot(
-                value=None,
-                label="Click to select (x, y) coordinates"
-            )
-            
-            # Initialize the coordinate selector
-            def create_selector():
-                import plotly.graph_objects as go
-                fig = go.Figure()
-                
-                # Add a square shape
-                fig.add_trace(go.Scatter(
-                    x=[0, 100, 100, 0, 0],
-                    y=[0, 0, 100, 100, 0],
-                    mode='lines',
-                    line=dict(color='black'),
-                    showlegend=False
-                ))
-                
-                # Update layout
-                fig.update_layout(
-                    width=300,
-                    height=300,
-                    margin=dict(l=0, r=0, t=0, b=0),
-                    xaxis=dict(
-                        range=[-5, 105],
-                        showgrid=False,
-                        zeroline=False,
-                        visible=False
-                    ),
-                    yaxis=dict(
-                        range=[-5, 105],
-                        showgrid=False,
-                        zeroline=False,
-                        visible=False,
-                        scaleanchor='x'
-                    ),
-                    plot_bgcolor='white'
-                )
-                return fig
-            
-            # Initialize the coordinate selector
-            coord_selector.value = create_selector()
-        
-        # Right column: Generated image
-        with gr.Column(scale=1):
-            output_image = gr.Image(
-                label="Generated Output",
-                height=300
-            )
-    
-    # Handle click events
-    coord_selector.select(
-        process_click,
-        inputs=[image_idx],
-        outputs=output_image
+    gr.Markdown("## Click Coordinate Detector\nClick anywhere on the white square to see coordinates")
+
+    # Display the white square
+    image = gr.Image(
+        label="Click on the white square",
+        value=create_white_square(),
+        interactive=True,
+        height=400,
+        width=400,
+        mirror_webcam=False,
     )
 
+    # Display coordinates
+    output_text = gr.Textbox(label="Coordinates")
+    print("oh yeah")
+    # Handle click events
+    image.select(get_click_coordinates, inputs=[], outputs=output_text)
+
 # Launch the app
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

inference.py

@@ -0,0 +1,96 @@
+import os
+import sys
+
+sys.path.append("..")
+
+from PIL import Image
+import matplotlib.pyplot as plt
+import hydra
+import omegaconf
+import jax
+import jax.numpy as jnp
+import optax
+from flax.training.train_state import TrainState
+from flax.serialization import from_bytes
+from huggingface_hub import snapshot_download
+
+# lpn imports
+from src.models.lpn import LPN
+from src.models.transformer import EncoderTransformer, DecoderTransformer
+from src.visualization import display_grid, ax_to_pil
+
+from utils import patch_target
+
+
+checkpoint_name = "quiet-thunder-789--checkpoint:v0"
+BLUE_LOCATION_INPUTS = {1: 13, 2: 9}
+
+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:
+    cfg = omegaconf.OmegaConf.load(f)
+patch_target(cfg)
+
+encoder = EncoderTransformer(hydra.utils.instantiate(cfg.encoder_transformer))
+decoder = DecoderTransformer(hydra.utils.instantiate(cfg.decoder_transformer))
+lpn = LPN(encoder=encoder, decoder=decoder)
+
+key = jax.random.PRNGKey(0)
+grids = jax.random.randint(
+    key,
+    (1, 3, decoder.config.max_rows, decoder.config.max_cols, 2),
+    minval=0,
+    maxval=decoder.config.vocab_size,
+)
+shapes = jax.random.randint(
+    key,
+    (1, 3, 2, 2),
+    minval=1,
+    maxval=min(decoder.config.max_rows, decoder.config.max_cols) + 1,
+)
+variables = lpn.init(
+    key, grids, shapes, dropout_eval=False, prior_kl_coeff=0.0, pairwise_kl_coeff=0.0, mode="mean"
+)
+learning_rate, linear_warmup_steps = 0, 0
+linear_warmup_scheduler = optax.warmup_exponential_decay_schedule(
+    init_value=learning_rate / (linear_warmup_steps + 1),
+    peak_value=learning_rate,
+    warmup_steps=linear_warmup_steps,
+    transition_steps=1,
+    end_value=learning_rate,
+    decay_rate=1.0,
+)
+optimizer = optax.chain(optax.clip_by_global_norm(1.0), optax.adamw(linear_warmup_scheduler))
+optimizer = optax.MultiSteps(optimizer, every_k_schedule=1)
+train_state = TrainState.create(apply_fn=lpn.apply, tx=optimizer, params=variables["params"])
+
+with open(os.path.join(local_dir, checkpoint_name, "state.msgpack"), "rb") as data_file:
+    byte_data = data_file.read()
+loaded_state = from_bytes(train_state, byte_data)
+
+generate_output_from_context = jax.jit(
+    lambda context, input, input_grid_shape: lpn.apply(
+        {"params": loaded_state.params},
+        context=context,
+        input=input,
+        input_grid_shape=input_grid_shape,
+        dropout_eval=True,
+        method=lpn._generate_output_from_context,
+    )
+)
+
+
+def generate_image(image_idx: int, x: float, y: float, eps: float = 1e-4) -> Image.Image:
+    # Create the input image
+    input = jnp.zeros(16, int).at[BLUE_LOCATION_INPUTS[image_idx]].set(1).reshape(4, 4)
+    # Ensure x and y are in [eps, 1 - eps]
+    x = min(1 - eps, max(eps, x))
+    y = min(1 - eps, max(eps, y))
+    # Convert x and y to context in R^2
+    context = jax.scipy.stats.norm.ppf(jnp.array([x, y]))
+    output_grids, _ = generate_output_from_context(
+        context=context[None], input=input[None], input_grid_shape=jnp.array([4, 4])[None]
+    )
+    output_grid = output_grids[0]
+    _, ax = plt.subplots(1, 1, figsize=(4, 4))
+    display_grid(ax=ax, grid=output_grid, grid_shape=jnp.array([4, 4]))
+    return ax_to_pil(ax)

requirements.txt

@@ -1,3 +1,3 @@
 gradio
 plotly
-git+https://github.com/clement-bonnet/lpn.git@edbe4722340719cc36b5a755fec7213cb8efb9f7
+git+https://github.com/clement-bonnet/lpn.git@f1bb82598454e897b3d4cb9f313d941943382877

utils.py

@@ -0,0 +1,12 @@
+import omegaconf
+
+
+def patch_target(config):
+    """Update the _target_ of cfg from src_v2 to src"""
+    for key, value in config.items():
+        if isinstance(value, omegaconf.DictConfig):
+            # Recursive call if the value is another DictConfig
+            patch_target(value)
+        elif isinstance(value, str) and value.startswith("src_v2"):
+            # Update the value if it matches the old_value
+            config[key] = value.replace("src_v2", "src")