Setup ZeroGPU

README.md

@@ -4,7 +4,7 @@ emoji: 💬
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 3.35.2
+sdk_version: 4.16.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -332,7 +332,7 @@ This is the demo for Dr-LLaVA. So far it could only be used for H&E stained Bone
 </ul>
                             """)
                 # Replace 'path_to_image' with the path to your image file
-                gr.Image(value="https://i.postimg.cc/tJzyq5Dh/Dr-LLa-VA-Fig-1.png", 
+                gr.Image(value="https://davidday.tw/wp-content/uploads/2024/08/Dr-LLa-VA-Fig-1.jpg", 
                      width=600, interactive=False, type="pil")
             with gr.Column(scale=3):
                 with gr.Row(elem_id="model_selector_row"):
@@ -497,7 +497,7 @@ def start_worker():
     ]
     return subprocess.Popen(worker_command)
 
-def get_args():
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int)
@@ -510,35 +510,22 @@ def get_args():
     parser.add_argument("--moderate", action="store_true")
     parser.add_argument("--embed", action="store_true")
     args = parser.parse_args()
-
-    return args
-
-
-def start_demo(args):
-    demo = build_demo(args.embed)
-    demo.queue(
-        concurrency_count=args.concurrency_count, status_update_rate=10, api_open=False
-    ).launch(server_name=args.host, server_port=args.port, share=args.share)
-
-
-if __name__ == "__main__":
-    args = get_args()
     logger.info(f"args: {args}")
 
     controller_proc = start_controller()
     worker_proc = start_worker()
 
     # Wait for worker and controller to start
-    time.sleep(10)
+    time.sleep(60)
 
-    exit_status = 0
-    try:
-        start_demo(args)
-    except Exception as e:
-        print(e)
-        exit_status = 1
-    finally:
-        worker_proc.kill()
-        controller_proc.kill()
-
-        sys.exit(exit_status)
+    models = get_model_list()
+
+    logger.info(args)
+    demo = build_demo(args.embed, concurrency_count=args.concurrency_count)
+    demo.queue(
+        api_open=False
+    ).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share
+    )

model_builder.py

@@ -23,7 +23,7 @@ from llava.model import *
 from llava.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
 
 
-def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", load_bf16=False):
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="cpu", load_bf16=False):
     kwargs = {"device_map": device_map}
 
     if load_8bit:

model_worker.py

@@ -14,6 +14,7 @@ import requests
 import torch
 import uvicorn
 from functools import partial
+import spaces
 
 from peft import PeftModel
 
@@ -72,6 +73,8 @@ class ModelWorker:
             self.model = PeftModel.from_pretrained(
                 self.model,
                 lora_path,
+                torch_device='cpu',
+                device_map="cpu",
             )
 
         if not no_register:
@@ -127,9 +130,10 @@ class ModelWorker:
             "queue_length": self.get_queue_length(),
         }
 
-    @torch.inference_mode()
+    @spaces.GPU
     def generate_stream(self, params):
         tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
+        logger.info(f'Model devices: {self.model.device}')
 
         prompt = params["prompt"]
         ori_prompt = prompt

requirements.txt

@@ -2,9 +2,7 @@
 tokenizers>=0.12.1
 torch==2.0.1
 torchvision==0.15.2
-deepspeed==0.9.5
-pydantic<2.0.0
-peft==0.4.0
+peft
 transformers==4.31.0
 accelerate==0.21.0
 bitsandbytes==0.41.0
@@ -12,5 +10,5 @@ sentencepiece==0.1.99
 einops==0.6.1
 einops-exts==0.0.4
 timm==0.6.13
-numpy<2
+httpx==0.24.0
 scipy