huaweilin/VTBench

1import os2import spaces3import gradio as gr4from src.data_processing import pil_to_tensor, tensor_to_pil5from PIL import Image6from src.model_processing import get_model7from huggingface_hub import snapshot_download8import torch9 10device = torch.device("cuda" if torch.cuda.is_available() else "cpu")11print(f"Running on: {device}")12 13MODEL_DIR = "./VTBench_models"14# if not os.path.exists(MODEL_DIR):15#     print("Downloading VTBench_models from Hugging Face...")16#     snapshot_download(17#         repo_id="huaweilin/VTBench_models",18#         local_dir=MODEL_DIR,19#         local_dir_use_symlinks=False20#     )21#     print("Download complete.")22 23example_image_paths = [f"assets/app_examples/{i}.png" for i in range(0, 5)]24 25model_name_mapping = {26    "SD3.5L": "SD3.5L",27    "chameleon": "Chameleon",28    # "flowmo_lo": "FlowMo Lo",29    # "flowmo_hi": "FlowMo Hi",30    # "gpt4o": "GPT-4o",31    "janus_pro_1b": "Janus Pro 1B/7B",32    "llamagen-ds8": "LlamaGen ds8",33    "llamagen-ds16": "LlamaGen ds16",34    "llamagen-ds16-t2i": "LlamaGen ds16 T2I",35    "maskbit_16bit": "MaskBiT 16bit",36    "maskbit_18bit": "MaskBiT 18bit",37    "open_magvit2": "OpenMagViT",38    "titok_b64": "Titok-b64",39    "titok_bl64": "Titok-bl64",40    "titok_s128": "Titok-s128",41    "titok_bl128": "Titok-bl128",42    "titok_l32": "Titok-l32",43    "titok_sl256": "Titok-sl256",44    "var_256": "VAR-256",45    "var_512": "VAR-512",46    "FLUX.1-dev": "FLUX.1-dev",47    "infinity_d32": "Infinity-d32",48    "infinity_d64": "Infinity-d64",49    "bsqvit": "BSQ-VIT",50}51 52display_to_internal = {v: k for k, v in model_name_mapping.items()}53 54def load_model(model_name):55    model, data_params = get_model(MODEL_DIR, model_name)56    model = model.to(device)57    model.eval()58    return model, data_params59 60# model_dict = {61#     model_name: load_model(model_name)62#     for model_name in model_name_mapping63# }64 65placeholder_image = Image.new("RGBA", (512, 512), (0, 0, 0, 0))66 67@spaces.GPU68def process_selected_models(uploaded_image, selected_display_names):69    if uploaded_image is None:70        return [gr.update(value="⚠️  Please upload an image before processing.", visible=True)] + \71               [gr.update() for _ in model_name_mapping]72 73    if not selected_display_names:74        return [gr.update(value="⚠️  Please select at least one model.", visible=True)] + \75               [gr.update() for _ in model_name_mapping]76 77    selected_results = []78    placeholder_results = []79 80    selected_internal = [display_to_internal[d] for d in selected_display_names]81 82    for model_name in model_name_mapping:83        label = model_name_mapping[model_name]84 85        if model_name in selected_internal:86            try:87                model, data_params = load_model(model_name)88                pixel_values = pil_to_tensor(uploaded_image, **data_params).unsqueeze(0).to(device)89                with torch.no_grad():90                    output = model(pixel_values)[0]91                reconstructed_image = tensor_to_pil(output[0].cpu(), **data_params)92 93                del model, pixel_values, output94                torch.cuda.empty_cache()95 96                result = gr.update(value=reconstructed_image, label=label)97            except Exception as e:98                print(f"Error in model {model_name}: {e}")99                result = gr.update(value=placeholder_image, label=f"{label} (Error)")100            selected_results.append(result)101        else:102            result = gr.update(value=placeholder_image, label=f"{label} (Not selected)")103            placeholder_results.append(result)104 105    return [gr.update(visible=False)] + selected_results + placeholder_results106 107 108with gr.Blocks() as demo:109    gr.Markdown("## VTBench")110    gr.Markdown("---")111 112    gr.Markdown("<span style='color:red; font-weight: bold;'>⚠️ ⚠️ ⚠️  If you encounter any errors, please try again — it usually works on the second attempt.</span>")113    gr.Markdown("👋 **Welcome to VTBench!** Upload an image, select models, and click 'Start Processing' to compare results side by side.")114    gr.Markdown("🔗 **Check out our GitHub repo:** [https://github.com/huawei-lin/VTBench](https://github.com/huawei-lin/VTBench)")115    with gr.Accordion("📘 Full Instructions", open=False):116        gr.Markdown("""117**VTBench User Guide**118 119- **Upload an image** or click one of the example images.120- **Select one or more models** from the list.121- Click **Start Processing** to run inference.122- Selected model outputs appear first, others show placeholders.123 124⚠️  *Each model is downloaded on first use. Please wait patiently the first time you run a model.*125""")126 127    image_input = gr.Image(128        type="pil",129        label="Upload an image",130        width=512,131        height=512,132    )133 134    gr.Markdown("### Click on an example image to use it as input:")135    example_rows = [example_image_paths[i:i+5] for i in range(0, len(example_image_paths), 5)]136    for row in example_rows:137        with gr.Row():138            for path in row:139                ex_img = gr.Image(140                    value=path,141                    show_label=False,142                    interactive=True,143                    width=256,144                    height=256,145                )146 147                def make_loader(p=path):148                    def load_img():149                        return Image.open(p)150                    return load_img151 152                ex_img.select(fn=make_loader(), outputs=image_input)153 154    gr.Markdown("---")155    gr.Markdown("⚠️  **The more models you select, the longer the processing time will be.**")156    gr.Markdown("*Note: Each model is downloaded on first use. Subsequent uses will load from cache and run faster.*")157 158    display_names = list(model_name_mapping.values())159    default_selected = ["SD3.5L", "Chameleon", "Janus Pro 1B/7B"]160 161    model_selector = gr.CheckboxGroup(162        choices=display_names,163        label="Select models to run",164        value=default_selected,165        interactive=True,166    )167 168    status_output = gr.Markdown("", visible=False)169    run_button = gr.Button("Start Processing")170 171    image_outputs = []172    model_names_ordered = list(model_name_mapping.keys())173    n_columns = 5174    output_rows = [model_names_ordered[i:i+n_columns] for i in range(0, len(model_names_ordered), n_columns)]175 176    with gr.Column():177        for row in output_rows:178            with gr.Row():179                for model_name in row:180                    display_name = model_name_mapping[model_name]181                    out_img = gr.Image(182                        label=f"{display_name} (Not run)",183                        value=placeholder_image,184                        width=512,185                        height=512,186                    )187                    image_outputs.append(out_img)188 189 190    run_button.click(191        fn=process_selected_models,192        inputs=[image_input, model_selector],193        outputs=[status_output] + image_outputs194    )195 196demo.launch()197 198