# Open Source Model Licensed under the Apache License Version 2.0
# and Other Licenses of the Third-Party Components therein:
# The below Model in this distribution may have been modified by THL A29 Limited
# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.

import os
import torch
from PIL import Image
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
import io
import tempfile

from hy3dgen.rembg import BackgroundRemover
from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline, FaceReducer, FloaterRemover, DegenerateFaceRemover
from hy3dgen.text2image import HunyuanDiTPipeline

app = FastAPI(title="Hunyuan3D-2 API")

# Initialize models at startup
rembg = BackgroundRemover()
model_path = 'tencent/Hunyuan3D-2'  # Use Hugging Face model path
t2i = HunyuanDiTPipeline('Tencent-Hunyuan--HunyuanDiT-v1.1-Diffusers-Distilled')
i23d = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(model_path)

class TextToImageRequest(BaseModel):
    prompt: str

@app.post("/image-to-3d/")
async def image_to_3d(file: UploadFile = File(...)):
    """Convert an image to 3D model"""
    try:
        # Read and process image
        content = await file.read()
        image = Image.open(io.BytesIO(content))
        image = image.resize((1024, 1024))
        
        if image.mode == 'RGB':
            image = rembg(image)

        # Generate mesh
        mesh = i23d(image=image, num_inference_steps=30, mc_algo='mc',
                   generator=torch.manual_seed(2025))[0]
        mesh = FloaterRemover()(mesh)
        mesh = DegenerateFaceRemover()(mesh)
        mesh = FaceReducer()(mesh)
        
        # Save mesh to temporary file
        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
            mesh.export(tmp.name)
            return FileResponse(tmp.name, media_type='model/gltf-binary', 
                             filename='model.glb')
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/text-to-3d/")
async def text_to_3d(request: TextToImageRequest):
    """Convert text prompt to 3D model"""
    try:
        # Generate image from text
        image = t2i(request.prompt)
        image = rembg(image)
        
        # Generate mesh
        mesh = i23d(image, num_inference_steps=30, mc_algo='mc')[0]
        mesh = FloaterRemover()(mesh)
        mesh = DegenerateFaceRemover()(mesh)
        mesh = FaceReducer()(mesh)
        
        # Save mesh to temporary file
        with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp:
            mesh.export(tmp.name)
            return FileResponse(tmp.name, media_type='model/gltf-binary',
                             filename='model.glb')
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)