import torch
import torch_npu
import argparse
from openmind import pipeline, is_torch_npu_available
from openmind import AutoImageProcessor
from openmind import AutoModel
from PIL import Image
import requests

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default=None,
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    if args.model_name_or_path:
        model_path = args.model_name_or_path
    else:
        model_path = snapshot_download(
            "GuangxiAICC/swin-small-finetuned-cifar100",
            revision="main",
            ignore_patterns=["*.h5", "*.ot", "*.msgpack"],
        )
    
    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
    
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    image = Image.open(requests.get(url, stream=True).raw)
    processor = AutoImageProcessor.from_pretrained(model_path)
    model = AutoModel.from_pretrained(model_path).to(device)

    inputs = processor(images=image, return_tensors="pt").to(device)
    outputs = model(**inputs)
    print("Predicted class:", outputs)

if __name__=="__main__":
    main()