fix: fix bugs.

gradio_app.py

@@ -333,22 +333,15 @@ def get_gradio_demo(model, tokenizer, image_processor) -> gr.Interface:
 
 
 def main(model_id: str = 'jxu124/TiO', host: str = "0.0.0.0", port: int = None):
-    if torch.cuda.is_available():
-        model = AutoModel.from_pretrained(
-            model_id,
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            device_map='cuda',
-            # load_in_4bit=True,
-            # bnb_4bit_compute_dtype=torch.float16,
-        )
-    else:
-        model = None
+    assert torch.cuda.is_available()
+    model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float16).cuda()
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
     image_processor = AutoImageProcessor.from_pretrained(model_id)
+
+    # ---- gradio demo ----
     model.get_gradio_demo(tokenizer, image_processor).queue(max_size=20).launch(server_name=host, server_port=port)
 
 
 if __name__ == "__main__":
     import fire
-    fire.Fire(main)
+    fire.Fire(main)

modeling_tio.py

@@ -103,7 +103,7 @@ def make_image_bucket_position(bucket_size, num_relative_distance):
     """
     coords_h = torch.arange(bucket_size)
     coords_w = torch.arange(bucket_size)
-    coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+    coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # 2, Wh, Ww
     coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
     relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
     relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2