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()