Spaces:
Running
on
A100
Running
on
A100
File size: 3,058 Bytes
174ae06 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# Copyright (c) 2025 NVIDIA CORPORATION.
# Licensed under the MIT license.
# Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license.
# LICENSE is in incl_licenses directory.
import argparse
import importlib.util
import json
import os
from pydantic import BaseModel
from termcolor import colored
import llava
from llava import conversation as clib
from llava.media import Image, Video, Sound
from llava.model.configuration_llava import JsonSchemaResponseFormat, ResponseFormat
from peft import PeftModel
import torch
def get_schema_from_python_path(path: str) -> str:
schema_path = os.path.abspath(path)
spec = importlib.util.spec_from_file_location("schema_module", schema_path)
schema_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(schema_module)
# Get the Main class from the loaded module
Main = schema_module.Main
assert issubclass(
Main, BaseModel
), f"The provided python file {path} does not contain a class Main that describes a JSON schema"
return Main.schema_json()
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--model-base", "-mb", type=str, required=True)
parser.add_argument("--model-path", "-mp", type=str, required=True)
parser.add_argument("--conv-mode", "-c", type=str, default="auto")
parser.add_argument("--text", type=str)
parser.add_argument("--media", type=str, nargs="+")
parser.add_argument("--json-mode", action="store_true")
parser.add_argument("--peft-mode", action="store_true")
parser.add_argument("--json-schema", type=str, default=None)
args = parser.parse_args()
# Convert json mode to response format
if not args.json_mode:
response_format = None
elif args.json_schema is None:
response_format = ResponseFormat(type="json_object")
else:
schema_str = get_schema_from_python_path(args.json_schema)
print(schema_str)
response_format = ResponseFormat(type="json_schema", json_schema=JsonSchemaResponseFormat(schema=schema_str))
# Load model
model = llava.load(args.model_base)
if args.peft_mode:
model = PeftModel.from_pretrained(
model,
args.model_path,
device_map="auto",
torch_dtype=torch.float16,
)
# Set conversation mode
clib.default_conversation = clib.conv_templates[args.conv_mode].copy()
# Prepare multi-modal prompt
prompt = []
if args.media is not None:
for media in args.media or []:
if any(media.endswith(ext) for ext in [".wav",".mp3", ".flac"]):
media = Sound(media)
else:
raise ValueError(f"Unsupported media type: {media}")
prompt.append(media)
if args.text is not None:
prompt.append(args.text)
# Generate response
response = model.generate_content(prompt, response_format=response_format)
print(colored(response, "cyan", attrs=["bold"]))
if __name__ == "__main__":
main()
|