Spaces:
Running
Running
| import io | |
| from typing import Optional | |
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| from transformers.utils.processor_visualizer_utils import ImageVisualizer | |
| MODELS = [ | |
| "openai/clip-vit-base-patch32", | |
| "HuggingFaceM4/Idefics3-8B-Llama3", | |
| ] | |
| def _fig_to_pil(fig) -> Image.Image: | |
| buf = io.BytesIO() | |
| fig.savefig(buf, format="png", bbox_inches="tight", dpi=160) | |
| buf.seek(0) | |
| return Image.open(buf).convert("RGB") | |
| def _run(model_id: str, image: Optional[Image.Image], use_sample: bool, add_grid: bool): | |
| viz = ImageVisualizer(model_id) | |
| captured = [] | |
| orig_show = plt.show | |
| def _capture_show(*_, **__): | |
| fig = plt.gcf() | |
| captured.append(fig) | |
| try: | |
| plt.show = _capture_show | |
| viz.visualize(images=None if use_sample else image, add_grid=add_grid) | |
| finally: | |
| plt.show = orig_show | |
| imgs = [_fig_to_pil(fig) for fig in captured] if captured else [] | |
| prompt_preview = viz.default_message(full_output=False) | |
| return imgs, prompt_preview | |
| def _resolve_and_run(model_pick, custom_model, image, use_sample, add_grid): | |
| model_id = (custom_model or "").strip() or (model_pick or "").strip() | |
| if not model_id: | |
| raise gr.Error("Pick a model or enter one.") | |
| return _run(model_id, image, use_sample, add_grid) | |
| def _on_image_change(_): | |
| return False # uncheck "use sample" when a custom image is set | |
| theme = gr.themes.Soft(primary_hue="orange", neutral_hue="gray") | |
| with gr.Blocks(title="Transformers Processor Visualizer", theme=theme) as demo: | |
| gr.Markdown("### Visualize what a processor feeds a vision–text model (uses the existing `ImageVisualizer`).") | |
| with gr.Row(): | |
| # LEFT: clickable models + custom field | |
| with gr.Column(scale=1, min_width=260): | |
| model_pick = gr.Radio( | |
| label="Models", | |
| choices=MODELS, | |
| value=MODELS[0], | |
| interactive=True, | |
| ) | |
| custom_model = gr.Textbox( | |
| label="Or type a model id", | |
| placeholder="owner/repo", | |
| lines=1, | |
| ) | |
| # RIGHT: controls + outputs | |
| with gr.Column(scale=3): | |
| with gr.Row(): | |
| add_grid = gr.Checkbox(label="Show patch grid", value=True) | |
| use_sample = gr.Checkbox(label="Use HF logo sample", value=True) | |
| image = gr.Image( | |
| label="Upload custom image", | |
| type="pil", | |
| height=140, | |
| sources=["upload"], | |
| ) | |
| gr.Markdown("### Render") | |
| gallery = gr.Gallery(label="Processor output", columns=[1], height=900) | |
| prompt = gr.Textbox(label="Compact chat template preview", lines=2) | |
| # Reactive updates: change model, toggle options, upload -> update immediately | |
| model_pick.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]) | |
| custom_model.submit(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]) | |
| add_grid.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]) | |
| use_sample.change(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]) | |
| image.change(_on_image_change, inputs=image, outputs=use_sample).then( | |
| _resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt] | |
| ) | |
| # Initial render so there is output before any interaction | |
| demo.load(_resolve_and_run, inputs=[model_pick, custom_model, image, use_sample, add_grid], outputs=[gallery, prompt]) | |
| if __name__ == "__main__": | |
| demo.launch() | |