File size: 690 Bytes
96fe658 |
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 |
import os
import torch
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['SWIFT_DEBUG'] = '1'
def test_deepseek_janus_pro_gene():
from swift.llm import infer_main, InferArguments
args = InferArguments(model='deepseek-ai/Janus-Pro-1B', infer_backend='pt')
infer_main(args)
def test_emu3_gen(infer_backend):
from swift.llm import infer_main, InferArguments
args = InferArguments(
model='BAAI/Emu3-Gen',
infer_backend=infer_backend,
stream=False,
use_chat_template=False,
top_k=2048,
max_new_tokens=40960)
infer_main(args)
if __name__ == '__main__':
# test_emu3_gen('pt')
test_deepseek_janus_pro_gene()
|