Update app.py

app.py

@@ -1,11 +1,24 @@
 # Copyright (c) Tencent Inc. All rights reserved.
+import time
 import os
+os.environ['PYTORCH_JIT'] = "0"
+os.system('mim install mmcv==2.0.1')
+# import spaces
 import sys
 import argparse
 import os.path as osp
 from io import BytesIO
 from functools import partial
 
+import spaces
+
+from mmengine.runner import Runner
+from mmengine.dataset import Compose
+from mmengine.runner.amp import autocast
+from mmengine.config import Config, DictAction, ConfigDict
+from mmdet.datasets import CocoDataset
+from mmyolo.registry import RUNNERS
+
 import cv2
 # import onnx
 import torch
@@ -15,12 +28,7 @@ import gradio as gr
 from PIL import Image
 import supervision as sv
 from torchvision.ops import nms
-from mmengine.runner import Runner
-from mmengine.dataset import Compose
-from mmengine.runner.amp import autocast
-from mmengine.config import Config, DictAction, ConfigDict
-from mmdet.datasets import CocoDataset
-from mmyolo.registry import RUNNERS
+
 
 from transformers import (AutoTokenizer, CLIPTextModelWithProjection)
 from transformers import (AutoProcessor, CLIPVisionModelWithProjection)
@@ -45,7 +53,7 @@ class LabelAnnotator(sv.LabelAnnotator):
 LABEL_ANNOTATOR = LabelAnnotator(text_padding=4,
                                  text_scale=0.5,
                                  text_thickness=1)
-
+@spaces.GPU
 def generate_image_embeddings(prompt_image,
                               vision_encoder,
                               vision_processor,
@@ -63,7 +71,7 @@ def generate_image_embeddings(prompt_image,
         img_feats = projector(img_feats)
     return img_feats
 
-
+@spaces.GPU
 def run_image(runner,
               vision_encoder,
               vision_processor,
@@ -254,6 +262,7 @@ if __name__ == '__main__':
     checkpoint = "weights/yolo_world_v2_l_image_prompt_adapter-719a7afb.pth"
     # load config
     cfg = Config.fromfile(config)
+    cfg.compile = False
     if cfg.get('work_dir', None) is None:
         cfg.work_dir = osp.join('./work_dirs',
                                 osp.splitext(osp.basename(config))[0])
@@ -264,7 +273,7 @@ if __name__ == '__main__':
         runner = Runner.from_cfg(cfg)
     else:
         runner = RUNNERS.build(cfg)
-
+    # runner.test()
     runner.call_hook('before_run')
     runner.load_or_resume()
     pipeline = cfg.test_dataloader.dataset.pipeline
@@ -276,7 +285,7 @@ if __name__ == '__main__':
     clip_model = "openai/clip-vit-base-patch32"
     vision_model = CLIPVisionModelWithProjection.from_pretrained(clip_model)
     processor = AutoProcessor.from_pretrained(clip_model)
-    device = 'cuda:0'
+    device = 'cuda'
     vision_model.to(device)
 
     texts = [' ']