added api_server

api_server.py

@@ -0,0 +1,476 @@
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
+# except for the third-party components listed below.
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined
+# in the repsective licenses of these third-party components.
+# Users must comply with all terms and conditions of original licenses of these third-party
+# components and must ensure that the usage of the third party components adheres to
+# all relevant laws and regulations.
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and
+# their software and algorithms, including trained model weights, parameters (including
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
+# fine-tuning enabling code and other elements of the foregoing made publicly available
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+"""
+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
+import time
+from io import BytesIO
+
+# Apply torchvision compatibility fix before other imports
+sys.path.insert(0, './hy3dshape')
+sys.path.insert(0, './hy3dpaint')
+
+try:
+    from torchvision_fix import apply_fix
+    apply_fix()
+except ImportError:
+    print("Warning: torchvision_fix module not found, proceeding without compatibility fix")
+except Exception as e:
+    print(f"Warning: Failed to apply torchvision fix: {e}")
+
+import torch
+import trimesh
+import uvicorn
+from PIL import Image
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, FileResponse
+
+# Updated imports to match gradio_app.py
+from hy3dshape import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, \
+    Hunyuan3DDiTFlowMatchingPipeline
+from hy3dshape.pipelines import export_to_trimesh
+from hy3dshape.rembg import BackgroundRemover
+from hy3dshape.utils import logger
+
+# Texture generation imports
+try:
+    from hy3dpaint.textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig
+    from hy3dpaint.convert_utils import create_glb_with_pbr_materials
+    HAS_TEXTUREGEN = True
+except ImportError:
+    print("Warning: Texture generation not available")
+    HAS_TEXTUREGEN = False
+
+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)))
+
+
+def export_mesh(mesh, save_folder, textured=False, type='glb'):
+    """
+    Export a mesh to a file in the specified folder, optionally including textures.
+
+    Args:
+        mesh (trimesh.Trimesh): The mesh object to export.
+        save_folder (str): Directory path where the mesh file will be saved.
+        textured (bool, optional): Whether to include textures/normals in the export. Defaults to False.
+        type (str, optional): File format to export ('glb' or 'obj' supported). Defaults to 'glb'.
+
+    Returns:
+        str: The full path to the exported mesh file.
+    """
+    if textured:
+        path = os.path.join(save_folder, f'textured_mesh.{type}')
+    else:
+        path = os.path.join(save_folder, f'white_mesh.{type}')
+    if type not in ['glb', 'obj']:
+        mesh.export(path)
+    else:
+        mesh.export(path, include_normals=textured)
+    return path
+
+
+class ModelWorker:
+    def __init__(self,
+                 model_path='tencent/Hunyuan3D-2.1',
+                 tex_model_path='tencent/Hunyuan3D-2.1',
+                 subfolder='hunyuan3d-dit-v2-1',
+                 device='cuda',
+                 enable_tex=False,
+                 low_vram_mode=False):
+        self.model_path = model_path
+        self.worker_id = worker_id
+        self.device = device
+        self.low_vram_mode = low_vram_mode
+        logger.info(f"Loading the model {model_path} on worker {worker_id} ...")
+
+        # Initialize background remover
+        self.rembg = BackgroundRemover()
+        
+        # Initialize shape generation pipeline
+        self.pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
+            model_path,
+            subfolder=subfolder,
+            use_safetensors=False,
+            device=device,
+        )
+            
+        # Initialize texture generation pipeline if enabled
+        if enable_tex and HAS_TEXTUREGEN:
+            try:
+                conf = Hunyuan3DPaintConfig(max_num_view=8, resolution=768)
+                conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth"
+                conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml"
+                conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr"
+                self.pipeline_tex = Hunyuan3DPaintPipeline(conf)
+            except Exception as e:
+                logger.error(f"Failed to initialize texture pipeline: {e}")
+                self.pipeline_tex = None
+        else:
+            self.pipeline_tex = None
+            
+        # Initialize mesh processing workers
+        self.floater_remove_worker = FloaterRemover()
+        self.degenerate_face_remove_worker = DegenerateFaceRemover()
+        self.face_reduce_worker = FaceReducer()
+
+    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):
+        start_time = time.time()
+        
+        # Handle input image
+        if 'image' in params:
+            image = params["image"]
+            image = load_image_from_base64(image)
+        else:
+            raise ValueError("No input image provided")
+
+        # Remove background if needed
+        if params.get('remove_background', True) or image.mode == "RGB":
+            image = self.rembg(image.convert('RGB'))
+
+        # Handle existing mesh or generate new one
+        if 'mesh' in params:
+            mesh = trimesh.load(BytesIO(base64.b64decode(params["mesh"])), file_type='glb')
+        else:
+            # Generate new mesh
+            seed = params.get("seed", 1234)
+            generator = torch.Generator(self.device).manual_seed(seed)
+            octree_resolution = params.get("octree_resolution", 256)
+            num_inference_steps = params.get("num_inference_steps", 5)
+            guidance_scale = params.get('guidance_scale', 5.0)
+            num_chunks = params.get('num_chunks', 8000)
+            
+            outputs = self.pipeline(
+                image=image,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                octree_resolution=octree_resolution,
+                num_chunks=num_chunks,
+                output_type='mesh'
+            )
+            
+            mesh = export_to_trimesh(outputs)[0]
+            logger.info("---Shape generation takes %s seconds ---" % (time.time() - start_time))
+
+        # Apply texture if requested
+        if params.get('texture', False) and self.pipeline_tex is not None:
+            # Post-process mesh for texture generation
+            mesh = self.floater_remove_worker(mesh)
+            mesh = self.degenerate_face_remove_worker(mesh)
+            mesh = self.face_reduce_worker(mesh, max_facenum=params.get('face_count', 40000))
+            
+            # Generate texture
+            tex_start_time = time.time()
+            temp_obj_path = os.path.join(SAVE_DIR, f'{str(uid)}_temp.obj')
+            mesh.export(temp_obj_path)
+            
+            text_path = os.path.join(SAVE_DIR, f'{str(uid)}_textured.obj')
+            self.pipeline_tex(mesh_path=temp_obj_path, 
+                             image_path=image, 
+                             output_mesh_path=text_path, 
+                             save_glb=False)
+            logger.info("---Texture generation takes %s seconds ---" % (time.time() - tex_start_time))
+            
+            # Convert to GLB with PBR materials if requested
+            file_type = params.get('type', 'glb')
+            if file_type == 'glb':
+                glb_path = os.path.join(SAVE_DIR, f'{str(uid)}.glb')
+                # Create texture paths (these would be generated by the texture pipeline)
+                textures = {
+                    'albedo': text_path.replace('.obj', '_albedo.png'),
+                    'metallic': text_path.replace('.obj', '_metallic.png'),
+                    'roughness': text_path.replace('.obj', '_roughness.jpg')
+                }
+                try:
+                    create_glb_with_pbr_materials(text_path, textures, glb_path)
+                    save_path = glb_path
+                except Exception as e:
+                    logger.warning(f"Failed to create PBR GLB, using regular export: {e}")
+                    mesh = trimesh.load(text_path)
+                    mesh.export(save_path)
+            else:
+                # Load textured mesh for other formats
+                mesh = trimesh.load(text_path)
+                mesh.export(save_path)
+        else:
+            # Export final mesh without texture
+            file_type = params.get('type', 'glb')
+            save_path = os.path.join(SAVE_DIR, f'{str(uid)}.{file_type}')
+            mesh.export(save_path)
+
+        if self.low_vram_mode:
+            torch.cuda.empty_cache()
+            
+        logger.info("---Total generation takes %s seconds ---" % (time.time() - start_time))
+        return save_path, uid
+
+
+app = FastAPI()
+from fastapi.middleware.cors import CORSMiddleware
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # 你可以指定允许的来源
+    allow_credentials=True,
+    allow_methods=["*"],  # 允许所有方法
+    allow_headers=["*"],  # 允许所有头部
+)
+
+
+@app.post("/generate")
+async def generate(request: Request):
+    logger.info("Worker generating...")
+    try:
+        params = await request.json()
+    except Exception as e:
+        logger.error(f"Failed to parse JSON request: {e}")
+        return JSONResponse({"error": "Invalid JSON request"}, status_code=400)
+    
+    # Validate required parameters
+    if not params.get('image'):
+        return JSONResponse({"error": "Image parameter is required"}, status_code=400)
+    
+    uid = uuid.uuid4()
+    try:
+        file_path, uid = worker.generate(uid, params)
+        return FileResponse(file_path)
+    except ValueError as e:
+        traceback.print_exc()
+        logger.error(f"Caught ValueError: {e}")
+        ret = {
+            "text": server_error_msg,
+            "error_code": 1,
+        }
+        return JSONResponse(ret, status_code=404)
+    except torch.cuda.CudaError as e:
+        logger.error(f"Caught torch.cuda.CudaError: {e}")
+        ret = {
+            "text": server_error_msg,
+            "error_code": 1,
+        }
+        return JSONResponse(ret, status_code=404)
+    except Exception as e:
+        logger.error(f"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...")
+    try:
+        params = await request.json()
+    except Exception as e:
+        logger.error(f"Failed to parse JSON request: {e}")
+        return JSONResponse({"error": "Invalid JSON request"}, status_code=400)
+    
+    # Validate required parameters
+    if not params.get('image'):
+        return JSONResponse({"error": "Image parameter is required"}, status_code=400)
+    
+    uid = uuid.uuid4()
+    try:
+        threading.Thread(target=worker.generate, args=(uid, params,)).start()
+        ret = {"uid": str(uid)}
+        return JSONResponse(ret, status_code=200)
+    except Exception as e:
+        logger.error(f"Failed to start generation thread: {e}")
+        ret = {"error": "Failed to start generation"}
+        return JSONResponse(ret, status_code=500)
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return JSONResponse({"status": "healthy", "worker_id": worker_id}, 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:
+        try:
+            base64_str = base64.b64encode(open(save_file_path, 'rb').read()).decode()
+            response = {'status': 'completed', 'model_base64': base64_str}
+            return JSONResponse(response, status_code=200)
+        except Exception as e:
+            logger.error(f"Error reading file {save_file_path}: {e}")
+            response = {'status': 'error', 'message': 'Failed to read generated file'}
+            return JSONResponse(response, status_code=500)
+
+
+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.1')
+    parser.add_argument("--subfolder", type=str, default='hunyuan3d-dit-v2-1')
+    parser.add_argument("--tex_model_path", type=str, default='tencent/Hunyuan3D-2.1')
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument('--enable_tex', action='store_true')
+    parser.add_argument('--low_vram_mode', action='store_true')
+    parser.add_argument('--cache-path', type=str, default='./gradio_cache')
+    parser.add_argument('--mc_algo', type=str, default='mc')
+    parser.add_argument('--compile', action='store_true')
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    # Update SAVE_DIR based on cache-path argument
+    SAVE_DIR = args.cache_path
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+    model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
+
+    worker = ModelWorker(
+        model_path=args.model_path, 
+        tex_model_path=args.tex_model_path,
+        subfolder=args.subfolder,
+        device=args.device, 
+        enable_tex=args.enable_tex,
+        low_vram_mode=args.low_vram_mode
+    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

docker/Dockerfile

@@ -53,8 +53,8 @@ RUN conda install -c conda-forge libstdcxx-ng -y
 
 RUN pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
 
-# Clone Hunyuan3D-2.1 repository
+# Clone Hunyuan3D-2.1 repository clone with out api_server 
-RUN git clone https://github.com/Tencent-Hunyuan/Hunyuan3D-2.1.git
+RUN git clone https://github.com/perfectproducts/Hunyuan3D-2.1.git
 
 # Install Python dependencies from modified requirements.txt
 RUN pip install -r Hunyuan3D-2.1/requirements.txt