File size: 3,809 Bytes
4c383ef | 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 | """ROCm/CUDA local WebUI for zero-shot vision transfer. No remote requests."""
import argparse,base64,io,json,sys,threading,traceback
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
from PIL import Image
from mugi_decision.vision import VisionDecisionEngine
ROOT=Path(__file__).resolve().parents[1]
def main():
ap=argparse.ArgumentParser()
ap.add_argument('--model',default=str(ROOT))
ap.add_argument('--adapter')
ap.add_argument('--port',type=int,default=7861)
ap.add_argument('--optimization',choices=['auto','none','hip','fast'],default='auto')
args=ap.parse_args()
engine=VisionDecisionEngine(model_path=args.model,adapter=args.adapter,optimization=args.optimization)
lock=threading.Lock()
class Handler(BaseHTTPRequestHandler):
def reply(self,status,data,kind='application/json; charset=utf-8'):
raw=json.dumps(data,ensure_ascii=False).encode('utf-8') if kind.startswith('application/json') else data
self.send_response(status)
self.send_header('Content-Type',kind)
self.send_header('Content-Length',str(len(raw)))
self.send_header('Cache-Control','no-store')
self.send_header('X-Content-Type-Options','nosniff')
self.end_headers();self.wfile.write(raw)
def do_GET(self):
if self.path=='/':self.reply(200,(ROOT/'webui/vision.html').read_bytes(),'text/html; charset=utf-8')
elif self.path=='/api/status':self.reply(200,{'ready':True,**engine.metadata})
else:self.reply(404,{'error':'Not found'})
def do_POST(self):
if self.path!='/api/predict':return self.reply(404,{'error':'Not found'})
origin=self.headers.get('Origin')
if origin and origin not in (f'http://127.0.0.1:{args.port}',f'http://localhost:{args.port}'):
return self.reply(403,{'error':'Local origin required'})
try:
size=int(self.headers.get('Content-Length','0'))
if not 0<size<=24*1024*1024:raise ValueError('请求过大（上限24MB）')
payload=json.loads(self.rfile.read(size))
picture=None
if payload.get('image'):
encoded=payload['image']
if not isinstance(encoded,str) or not encoded.startswith('data:image/'):
raise ValueError('需要图片 data URL')
picture=Image.open(io.BytesIO(base64.b64decode(encoded.split(',',1)[1],validate=True)))
if picture.width*picture.height>20_000_000:raise ValueError('图片不得超过2000万像素')
picture.load()
if not lock.acquire(blocking=False):return self.reply(409,{'error':'模型正在处理另一个请求，请稍后再试'})
try:
result=engine.predict(payload.get('description',''),payload.get('options',[]),picture,
int(payload.get('image_size',512)),int(payload.get('max_length',4096)),bool(payload.get('verify_cache',False)))
finally:lock.release()
self.reply(200,result)
except (ValueError,TypeError,KeyError,Image.DecompressionBombError) as exc:
self.reply(400,{'error':str(exc)})
except Exception as exc:
traceback.print_exc()
self.reply(500,{'error':f'{type(exc).__name__}: {exc}'})
server=ThreadingHTTPServer(('127.0.0.1',args.port),Handler)
print(json.dumps(engine.metadata,ensure_ascii=False),flush=True)
print(f'WebUI ready: http://127.0.0.1:{args.port}',flush=True)
server.serve_forever()
if __name__=='__main__':main()
|