Create api_server.py

api_server.py

@@ -0,0 +1,273 @@
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import base64
+import logging
+import logging.handlers
+import os
+import sys
+import tempfile
+import threading
+import traceback
+import uuid
+from io import BytesIO
+
+import torch
+import trimesh
+import uvicorn
+from PIL import Image
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, FileResponse
+
+from hy3dgen.rembg import BackgroundRemover
+from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline, FloaterRemover, DegenerateFaceRemover, FaceReducer
+from hy3dgen.texgen import Hunyuan3DPaintPipeline
+from hy3dgen.text2image import HunyuanDiTPipeline
+
+LOGDIR = '.'
+
+server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+
+handler = None
+
+
+def build_logger(logger_name, logger_filename):
+    global handler
+
+    formatter = logging.Formatter(
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    # Set the format of root handlers
+    if not logging.getLogger().handlers:
+        logging.basicConfig(level=logging.INFO)
+    logging.getLogger().handlers[0].setFormatter(formatter)
+
+    # Redirect stdout and stderr to loggers
+    stdout_logger = logging.getLogger("stdout")
+    stdout_logger.setLevel(logging.INFO)
+    sl = StreamToLogger(stdout_logger, logging.INFO)
+    sys.stdout = sl
+
+    stderr_logger = logging.getLogger("stderr")
+    stderr_logger.setLevel(logging.ERROR)
+    sl = StreamToLogger(stderr_logger, logging.ERROR)
+    sys.stderr = sl
+
+    # Get logger
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    # Add a file handler for all loggers
+    if handler is None:
+        os.makedirs(LOGDIR, exist_ok=True)
+        filename = os.path.join(LOGDIR, logger_filename)
+        handler = logging.handlers.TimedRotatingFileHandler(
+            filename, when='D', utc=True, encoding='UTF-8')
+        handler.setFormatter(formatter)
+
+        for name, item in logging.root.manager.loggerDict.items():
+            if isinstance(item, logging.Logger):
+                item.addHandler(handler)
+
+    return logger
+
+
+class StreamToLogger(object):
+    """
+    Fake file-like stream object that redirects writes to a logger instance.
+    """
+
+    def __init__(self, logger, log_level=logging.INFO):
+        self.terminal = sys.stdout
+        self.logger = logger
+        self.log_level = log_level
+        self.linebuf = ''
+
+    def __getattr__(self, attr):
+        return getattr(self.terminal, attr)
+
+    def write(self, buf):
+        temp_linebuf = self.linebuf + buf
+        self.linebuf = ''
+        for line in temp_linebuf.splitlines(True):
+            # From the io.TextIOWrapper docs:
+            #   On output, if newline is None, any '\n' characters written
+            #   are translated to the system default line separator.
+            # By default sys.stdout.write() expects '\n' newlines and then
+            # translates them so this is still cross platform.
+            if line[-1] == '\n':
+                self.logger.log(self.log_level, line.rstrip())
+            else:
+                self.linebuf += line
+
+    def flush(self):
+        if self.linebuf != '':
+            self.logger.log(self.log_level, self.linebuf.rstrip())
+        self.linebuf = ''
+
+
+def pretty_print_semaphore(semaphore):
+    if semaphore is None:
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"
+
+
+SAVE_DIR = 'gradio_cache'
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("controller", f"{SAVE_DIR}/controller.log")
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+class ModelWorker:
+    def __init__(self, model_path='tencent/Hunyuan3D-2', device='cuda'):
+        self.model_path = model_path
+        self.worker_id = worker_id
+        self.device = device
+        logger.info(f"Loading the model {model_path} on worker {worker_id} ...")
+
+        self.rembg = BackgroundRemover()
+        self.pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(model_path, device=device)
+        self.pipeline_t2i = HunyuanDiTPipeline('Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled',
+                                               device=device)
+        self.pipeline_tex = Hunyuan3DPaintPipeline.from_pretrained(model_path)
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return args.limit_model_concurrency - model_semaphore._value + (len(
+                model_semaphore._waiters) if model_semaphore._waiters is not None else 0)
+
+    def get_status(self):
+        return {
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate(self, uid, params):
+        if 'image' in params:
+            image = params["image"]
+            image = load_image_from_base64(image)
+        else:
+            if 'text' in params:
+                text = params["text"]
+                image = self.pipeline_t2i(text)
+            else:
+                raise ValueError("No input image or text provided")
+
+        image = self.rembg(image)
+        params['image'] = image
+
+        if 'mesh' in params:
+            mesh = trimesh.load(BytesIO(base64.b64decode(params["mesh"])), file_type='glb')
+        else:
+            seed = params.get("seed", 1234)
+            params['generator'] = torch.Generator(self.device).manual_seed(seed)
+            params['octree_resolution'] = params.get("octree_resolution", 256)
+            params['num_inference_steps'] = params.get("num_inference_steps", 30)
+            params['guidance_scale'] = params.get('guidance_scale', 7.5)
+            params['mc_algo'] = 'mc'
+            mesh = self.pipeline(**params)[0]
+
+        if params.get('texture', False):
+            mesh = FloaterRemover()(mesh)
+            mesh = DegenerateFaceRemover()(mesh)
+            mesh = FaceReducer()(mesh, max_facenum=params.get('face_count', 40000))
+            mesh = self.pipeline_tex(mesh, image)
+
+        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as temp_file:
+            mesh.export(temp_file.name)
+            mesh = trimesh.load(temp_file.name)
+            temp_file.close()
+            os.unlink(temp_file.name)
+            save_path = os.path.join(SAVE_DIR, f'{str(uid)}.glb')
+            mesh.export(save_path)
+
+        torch.cuda.empty_cache()
+        return save_path, uid
+
+
+app = FastAPI()
+
+
+@app.post("/generate")
+async def generate(request: Request):
+    logger.info("Worker generating...")
+    params = await request.json()
+    uid = uuid.uuid4()
+    try:
+        file_path, uid = worker.generate(uid, params)
+        return FileResponse(file_path)
+    except ValueError as e:
+        traceback.print_exc()
+        print("Caught ValueError:", e)
+        ret = {
+            "text": server_error_msg,
+            "error_code": 1,
+        }
+        return JSONResponse(ret, status_code=404)
+    except torch.cuda.CudaError as e:
+        print("Caught torch.cuda.CudaError:", e)
+        ret = {
+            "text": server_error_msg,
+            "error_code": 1,
+        }
+        return JSONResponse(ret, status_code=404)
+    except Exception as e:
+        print("Caught Unknown Error", e)
+        traceback.print_exc()
+        ret = {
+            "text": server_error_msg,
+            "error_code": 1,
+        }
+        return JSONResponse(ret, status_code=404)
+
+
+@app.post("/send")
+async def generate(request: Request):
+    logger.info("Worker send...")
+    params = await request.json()
+    uid = uuid.uuid4()
+    threading.Thread(target=worker.generate, args=(uid, params,)).start()
+    ret = {"uid": str(uid)}
+    return JSONResponse(ret, status_code=200)
+
+
+@app.get("/status/{uid}")
+async def status(uid: str):
+    save_file_path = os.path.join(SAVE_DIR, f'{uid}.glb')
+    print(save_file_path, os.path.exists(save_file_path))
+    if not os.path.exists(save_file_path):
+        response = {'status': 'processing'}
+        return JSONResponse(response, status_code=200)
+    else:
+        base64_str = base64.b64encode(open(save_file_path, 'rb').read()).decode()
+        response = {'status': 'completed', 'model_base64': base64_str}
+        return JSONResponse(response, status_code=200)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8081)
+    parser.add_argument("--model_path", type=str, default='tencent/Hunyuan3D-2')
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+
+    worker = ModelWorker(model_path=args.model_path, device=args.device)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")