File size: 1,404 Bytes
1e6a1e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
print("app Started")
import torch
from moondream2.config import MoondreamConfig
from moondream2.moondream import MoondreamModel
import torch.profiler
config = MoondreamConfig()
device = "cuda"
model = MoondreamModel(config, setup_caches=False).to(device)
from safetensors.torch import load_file
weights_path = "moondream2/model.safetensors" # Path to your local weights file
state_dict = load_file(weights_path, device=device)
new_state_dict = {}
for key, value in state_dict.items():
# Remove 'model.' prefix if it exists
if key.startswith('model.'):
new_key = key[6:] # Skip the first 6 characters ('model.')
else:
new_key = key
new_state_dict[new_key] = value
state_dict = new_state_dict
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=True)
model._setup_caches()
from PIL import Image
image = Image.open("example.png")
query = "home icon at the bottom"
warmup_iters = 2
for i in range(3):
if i == warmup_iters: torch.cuda.cudart().cudaProfilerStart()
if i >= warmup_iters: torch.cuda.nvtx.range_push("iteration{}".format(i))
if i >= warmup_iters: torch.cuda.nvtx.range_push("forward")
points = model.point(image, query)["points"]
if i >= warmup_iters: torch.cuda.nvtx.range_pop()
if i >= warmup_iters: torch.cuda.nvtx.range_pop()
torch.cuda.cudart().cudaProfilerStop()
|