mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 15:41:00 +02:00
Vladimir Mandic
51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
import time
|
|
import logging
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from .serve import OpenAIServer
|
|
from modules import logger
|
|
|
|
|
|
logger.setup_logging(debug=True)
|
|
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
|
logging.getLogger("httpx").setLevel(logging.WARNING)
|
|
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
|
logging.getLogger("asyncio").setLevel(logging.WARNING)
|
|
|
|
|
|
logger.log.info("OpenAI: load model...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"Qwen/Qwen3-0.6B",
|
|
device_map="auto",
|
|
dtype=torch.bfloat16,
|
|
trust_remote_code=True,
|
|
attn_implementation="sdpa",
|
|
# attn_implementation="eager",
|
|
# attn_implementation="flash_attention_2",
|
|
)
|
|
model.config.use_flash_attention = True
|
|
|
|
logger.log.info("OpenAI: eval model...")
|
|
model.eval()
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
"Qwen/Qwen3-0.6B",
|
|
trust_remote_code=True
|
|
)
|
|
|
|
logger.log.info("OpenAI: create server...")
|
|
server = OpenAIServer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
host="127.0.0.1",
|
|
port=8000
|
|
)
|
|
|
|
logger.log.info("OpenAI: start server...")
|
|
server.start()
|
|
while True:
|
|
try:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
server.stop()
|
|
break
|