Add upload-model and list-models subcommands to emcli

This commit is contained in:
Fam Zheng 2025-11-25 22:18:58 +00:00
parent eb68445134
commit ce280586fe

View File

@ -5,8 +5,10 @@ import sys
import json
import requests
import argparse
import logging
from multiprocessing import Pool
from tqdm import tqdm
import oss2
class SubCommand(object):
""" Base class of subcommand"""
@ -275,6 +277,81 @@ class UserInfoCommand(SubCommand):
r.raise_for_status()
print(r.json())
class UploadModelCommand(SubCommand):
name = "upload-model"
want_argv = False
help = "Upload model file to OSS bucket for qr-verify API"
def setup_args(self, parser):
parser.add_argument("file", help="Model file to upload")
def do(self, args, argv):
# OSS credentials (same as in emblem5/ai/ossclient.py)
oss_ak = 'LTAI5tC2qXGxwHZUZP7DoD1A'
oss_sk = 'qPo9O6ZvEfqo4t8oflGEm0DoxLHJhm'
# OSS endpoint and bucket (same as in emblem5/ai/server.py download_model)
endpoint = 'https://oss-rg-china-mainland.aliyuncs.com'
bucket_name = 'emblem-models'
model_file = args.file
if not os.path.exists(model_file):
print(f"Error: File not found: {model_file}", file=sys.stderr)
return 1
model_name = os.path.basename(model_file)
print(f"Uploading {model_file} to {bucket_name}/{model_name}...")
try:
auth = oss2.Auth(oss_ak, oss_sk)
bucket = oss2.Bucket(auth, endpoint, bucket_name)
bucket.put_object_from_file(model_name, model_file)
print(f"Successfully uploaded {model_name} to OSS bucket {bucket_name}")
print(f"The qr-verify API will be able to download it as: {model_name}")
except Exception as e:
print(f"Error uploading model: {e}", file=sys.stderr)
return 1
class ListModelsCommand(SubCommand):
name = "list-models"
want_argv = False
help = "List models available in OSS bucket for qr-verify API"
def setup_args(self, parser):
pass
def do(self, args, argv):
# OSS credentials (same as in emblem5/ai/ossclient.py)
oss_ak = 'LTAI5tC2qXGxwHZUZP7DoD1A'
oss_sk = 'qPo9O6ZvEfqo4t8oflGEm0DoxLHJhm'
# OSS endpoint and bucket (same as in emblem5/ai/server.py download_model)
endpoint = 'https://oss-rg-china-mainland.aliyuncs.com'
bucket_name = 'emblem-models'
try:
auth = oss2.Auth(oss_ak, oss_sk)
bucket = oss2.Bucket(auth, endpoint, bucket_name)
print(f"Models in OSS bucket {bucket_name}:")
print("-" * 60)
models = []
for obj in oss2.ObjectIterator(bucket):
models.append(obj.key)
if not models:
print("No models found in bucket.")
else:
models.sort()
for i, model_name in enumerate(models, 1):
print(f"{i:3d}. {model_name}")
print("-" * 60)
print(f"Total: {len(models)} model(s)")
except Exception as e:
print(f"Error listing models: {e}", file=sys.stderr)
return 1
def global_args(parser):
parser.add_argument("--env", "-E", help="Env", default="prod")
parser.add_argument("-D", "--debug", action="store_true",