add examples

app.py

@@ -3,6 +3,10 @@ import logging
 import sys
 from config import WEAVE_PROJECT, WANDB_API_KEY
 import weave
+from model_utils import get_model_summary, install_flash_attn
+
+# Install required package
+install_flash_attn()
 
 weave.init(WEAVE_PROJECT)
 
@@ -186,9 +190,15 @@ def detect_brain_tumor_florence2(image, seg_input, debug: bool = True):
 
     return (image_with_bboxes, annotations)
 
+def handle_model_summary(model_name):
+    model_summary, error_message = get_model_summary(model_name)
+    if error_message:
+        return error_message, ""
+    return model_summary, ""
+
 INTRO_TEXT="# 🔬🧠 OmniScience -- Agentic Imaging Analysis 🤖🧫"
 
-with gr.Blocks(css="style.css") as demo:
+with gr.Blocks(theme="sudeepshouche/minimalist") as demo:
     gr.Markdown(INTRO_TEXT)
     with gr.Tab("Object Detection - Owl V2"):
         with gr.Row():
@@ -296,7 +306,29 @@ with gr.Blocks(css="style.css") as demo:
             fn=detect_brain_tumor_dino,
             inputs=seg_inputs,
             outputs=seg_outputs,
-        )        
+        )
+
+    with gr.Tab("Model Explorer"):
+        gr.Markdown("## Retrieve and Display Model Architecture")
+        model_name_input = gr.Textbox(label="Model Name", placeholder="Enter the model name to retrieve its architecture...")
+        vision_examples = gr.Examples(
+            examples=[
+                ["facebook/sam-vit-huge"],
+                ["google/owlv2-base-patch16-ensemble"],
+                ["IDEA-Research/grounding-dino-base"],
+                ["microsoft/Florence-2-large-ft"],
+                ["google/paligemma-3b-mix-224"],
+                ["llava-hf/llava-v1.6-mistral-7b-hf"],
+                ["vikhyatk/moondream2"],
+                ["microsoft/Phi-3-vision-128k-instruct"],
+                ["HuggingFaceM4/idefics2-8b-chatty"]
+            ],
+            inputs=model_name_input
+        )
+        model_output = gr.Textbox(label="Model Architecture", lines=20, placeholder="Model architecture will appear here...", show_copy_button=True)
+        error_output = gr.Textbox(label="Error", lines=10, placeholder="Exceptions will appear here...", show_copy_button=True)
+        model_submit_button = gr.Button("Submit")
+        model_submit_button.click(fn=handle_model_summary, inputs=model_name_input, outputs=[model_output, error_output])
 
 if __name__ == "__main__":
     demo.queue(max_size=10).launch(debug=True)

model_utils.py

@@ -0,0 +1,66 @@
+import subprocess
+import os
+import torch
+from transformers import BitsAndBytesConfig, AutoConfig, AutoModelForCausalLM, LlavaNextForConditionalGeneration, LlavaForConditionalGeneration, PaliGemmaForConditionalGeneration, Idefics2ForConditionalGeneration, Owlv2ForObjectDetection, GroundingDinoForObjectDetection, SamModel
+import spaces
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+def install_flash_attn():
+    subprocess.run(
+        "pip install flash-attn --no-build-isolation",
+        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+        shell=True,
+    )
+
+ARCHITECTURE_MAP = {
+    "LlavaNextForConditionalGeneration": LlavaNextForConditionalGeneration,
+    "LlavaForConditionalGeneration": LlavaForConditionalGeneration,
+    "PaliGemmaForConditionalGeneration": PaliGemmaForConditionalGeneration,
+    "Idefics2ForConditionalGeneration": Idefics2ForConditionalGeneration,
+    "Owlv2ForObjectDetection": Owlv2ForObjectDetection,
+    "GroundingDinoForObjectDetection": GroundingDinoForObjectDetection,
+    "SamModel": SamModel,
+    "AutoModelForCausalLM": AutoModelForCausalLM
+}
+
+
+@spaces.GPU
+def get_model_summary(model_name):
+    try:
+        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+        architecture = config.architectures[0]
+        quantization_config = getattr(config, 'quantization_config', None)
+
+        if quantization_config:
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=quantization_config.get('load_in_4bit', False),
+                load_in_8bit=quantization_config.get('load_in_8bit', False),
+                bnb_4bit_compute_dtype=quantization_config.get('bnb_4bit_compute_dtype', torch.float16),
+                bnb_4bit_quant_type=quantization_config.get('bnb_4bit_quant_type', 'nf4'),
+                bnb_4bit_use_double_quant=quantization_config.get('bnb_4bit_use_double_quant', False),
+                llm_int8_enable_fp32_cpu_offload=quantization_config.get('llm_int8_enable_fp32_cpu_offload', False),
+                llm_int8_has_fp16_weight=quantization_config.get('llm_int8_has_fp16_weight', False),
+                llm_int8_skip_modules=quantization_config.get('llm_int8_skip_modules', None),
+                llm_int8_threshold=quantization_config.get('llm_int8_threshold', 6.0),
+            )
+        else:
+            bnb_config = None
+
+        model_class = ARCHITECTURE_MAP.get(architecture, AutoModelForCausalLM)
+        model = model_class.from_pretrained(
+            model_name, config=bnb_config, trust_remote_code=True
+        )
+
+        if model and not quantization_config:
+            model = model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+
+        model_summary = str(model) if model else "Model architecture not found."
+        config_content = config.to_json_string() if config else "Configuration not found."
+        return f"## Model Architecture\n\n{model_summary}\n\n## Configuration\n\n{config_content}", ""
+    except ValueError as ve:
+        return "", f"ValueError: {ve}"
+    except EnvironmentError as ee:
+        return "", f"EnvironmentError: {ee}"
+    except Exception as e:
+        return "", str(e)

requirements.txt

@@ -4,4 +4,6 @@ pillow
 pillow-heif
 weave
 huggingface-hub
-gradio
+gradio
+transformers
+spaces