Create app.py

app.py

@@ -0,0 +1,48 @@
+import gradio as gr
+import torch
+import numpy as np
+from transformers import AutoModel
+from theia.decoding import load_feature_stats, prepare_depth_decoder, prepare_mask_generator, decode_everything
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+def run_theia(image):
+    theia_model = AutoModel.from_pretrained("theaiinstitute/theia-base-patch16-224-cdiv", trust_remote_code=True)
+    theia_model = theia_model.to(device)
+    target_model_names = [
+        "google/vit-huge-patch14-224-in21k",
+        "facebook/dinov2-large",
+        "openai/clip-vit-large-patch14",
+        "facebook/sam-vit-huge",
+        "LiheYoung/depth-anything-large-hf",
+    ]
+    feature_means, feature_vars = load_feature_stats(target_model_names, stat_file_root="../../../feature_stats")
+    
+    mask_generator, sam_model = prepare_mask_generator(device)
+    depth_anything_model_name = "LiheYoung/depth-anything-large-hf"
+    depth_anything_decoder, _ = prepare_depth_decoder(depth_anything_model_name, device)
+    
+    images = [image]
+    
+    theia_decode_results, gt_decode_results = decode_everything(
+        theia_model=theia_model,
+        feature_means=feature_means,
+        feature_vars=feature_vars,
+        images=images,
+        mask_generator=mask_generator,
+        sam_model=sam_model,
+        depth_anything_decoder=depth_anything_decoder,
+        pred_iou_thresh=0.5,
+        stability_score_thresh=0.7,
+        gt=True,
+        device=device,
+    )
+    
+    vis_video = np.stack(
+        [np.vstack([tr, gtr]) for tr, gtr in zip(theia_decode_results, gt_decode_results, strict=False)]
+    )
+
+    return vis_video
+
+demo = gr.Interface(fn=run_theia, inputs="image", outputs="image")
+demo.launch()