add: vitb16 support and more user inputs.

app.py

@@ -132,7 +132,6 @@ def run(pil_image: Image.Image, epoch: str, weight_type: str, image_size: int):
     image_size = int(image_size)
     pending_img = Image.new("RGB", (image_size, image_size), color=(18, 18, 18))
 
-    # 6 total positions: ViT-S [dino, ibot, lejepa], ViT-B [dino, ibot, lejepa]
     results = [pending_img] * 6
     yield tuple(results)
 
@@ -144,7 +143,6 @@ def run(pil_image: Image.Image, epoch: str, weight_type: str, image_size: int):
         for model_key in MODEL_KEYS:
             repo_id = MODEL_IDS[arch][model_key]
 
-            # LeJEPA only supports student weights
             current_weight = "student" if model_key == "LeJEPA" else weight_type
             revision = f"{epoch}/{current_weight}"
 
@@ -153,7 +151,6 @@ def run(pil_image: Image.Image, epoch: str, weight_type: str, image_size: int):
                 results[idx] = pca_vis(model, image_tensor, image_size)
             except Exception as e:
                 print(f"Error processing {repo_id} ({revision}): {e}")
-                # Create an error placeholder card if a model/revision download fails
                 error_canvas = Image.new("RGB", (image_size, image_size), color=(40, 20, 20))
                 results[idx] = error_canvas
 
@@ -189,7 +186,6 @@ CSS = """
     color: #374151;
     padding: 0.25rem 0;
 }
-/* Ensure strict rigid layouts for outputs to avoid layout shifting */
 .output-col {
     display: flex !important;
     flex-direction: column !important;
@@ -204,14 +200,22 @@ CSS = """
     max-height: 350px !important;
     width: 100% !important;
 }
-.subtitle-row a, .model-label a {
+.subtitle-row a, .model-label a, .custom-footer a {
     color: inherit;
     text-decoration: underline;
     text-decoration-color: #d1d5db;
 }
-.model-label a:hover, .subtitle-row a:hover {
+.model-label a:hover, .subtitle-row a:hover, .custom-footer a:hover {
     text-decoration-color: currentColor;
 }
+.custom-footer {
+    text-align: center;
+    margin-top: 2.5rem;
+    padding-top: 1rem;
+    border-top: 1px solid #e5e7eb;
+    font-size: 0.8rem;
+    color: #9ca3af;
+}
 footer { display: none !important; }
 """
 
@@ -273,28 +277,43 @@ with gr.Blocks(css=CSS, title="SSL ViT PCA Visualiser") as demo:
             gr.HTML('<div class="arch-header">ViT-S/16 Grid</div>')
             with gr.Row(equal_height=True):
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">DiNO (S/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-S/16"]["DiNO"]}" target="_blank">DiNO (S/16)</a></div>')
                     out_dino_s = gr.Image(show_label=False, interactive=False)
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">iBOT (S/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-S/16"]["iBOT"]}" target="_blank">iBOT (S/16)</a></div>')
                     out_ibot_s = gr.Image(show_label=False, interactive=False)
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">LeJEPA (S/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-S/16"]["LeJEPA"]}" target="_blank">LeJEPA (S/16)</a></div>')
                     out_lejepa_s = gr.Image(show_label=False, interactive=False)
 
             # ── ViT-B/16 Row ──
             gr.HTML('<div class="arch-header">ViT-B/16 Grid</div>')
             with gr.Row(equal_height=True):
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">DiNO (B/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-B/16"]["DiNO"]}" target="_blank">DiNO (B/16)</a></div>')
                     out_dino_b = gr.Image(show_label=False, interactive=False)
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">iBOT (B/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-B/16"]["iBOT"]}" target="_blank">iBOT (B/16)</a></div>')
                     out_ibot_b = gr.Image(show_label=False, interactive=False)
                 with gr.Column(elem_classes="output-col"):
-                    gr.HTML('<div class="model-label">LeJEPA (B/16)</div>')
+                    gr.HTML(
+                        f'<div class="model-label"><a href="https://huggingface.co/{MODEL_IDS["ViT-B/16"]["LeJEPA"]}" target="_blank">LeJEPA (B/16)</a></div>')
                     out_lejepa_b = gr.Image(show_label=False, interactive=False)
 
+    # Custom Clean Footer layout containing links to organization and codebase
+    gr.HTML("""
+        <div class="custom-footer">
+            Models: <a href="https://huggingface.co/OK-AI" target="_blank">OK-AI on HuggingFace</a>
+            &nbsp;·&nbsp;
+            Code: <a href="https://github.com/Open-Knowledge-AI/lite_ssl" target="_blank">lite_ssl Github</a>
+        </div>
+    """)
+
     # Wire outputs orderly following the exact resolution pattern tracking inside the `run` loop
     output_targets = [
         out_dino_s, out_ibot_s, out_lejepa_s,