server.py

import os
import sys
import uuid
import torch
import numpy as np
import asyncio
from PIL import Image
from typing import Dict, Any, List
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse

# Add roots to python path
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
if BASE_DIR not in sys.path:
    sys.path.append(BASE_DIR)
TRELLIS_DIR = os.path.join(BASE_DIR, "TRELLIS-main")
if TRELLIS_DIR not in sys.path:
    sys.path.append(TRELLIS_DIR)

# Force xformers backend
os.environ['ATTN_BACKEND'] = 'xformers'
os.environ['SPCONV_ALGO'] = 'native'

from formaai import FormaAi

app = FastAPI(title="FormaAi Custom API Server")

# Enable CORS for convenience
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize FormaAi model
print("[Backend] Initializing FormaAi pipeline on GPU/CPU...")
device = "cuda" if torch.cuda.is_available() else "cpu"
model = FormaAi(device=device)

# Directory structure
OUTPUT_DIR = os.path.join(BASE_DIR, "static/outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# In-memory task database
tasks: Dict[str, Dict[str, Any]] = {}

def run_pipeline_task(task_id: str, image: Image.Image, params: Dict[str, Any]):
    try:
        tasks[task_id]["status"] = "processing"
        tasks[task_id]["progress"] = 10
        tasks[task_id]["stage"] = "Вырезание фона и центрирование (Rembg)..."
        
        task_dir = os.path.join(OUTPUT_DIR, task_id)
        os.makedirs(task_dir, exist_ok=True)
        
        # 1. Preprocess input image
        processed_image = model.preprocess(image)
        preprocessed_path = os.path.join(task_dir, "preprocessed.png")
        processed_image.save(preprocessed_path)
        
        tasks[task_id]["progress"] = 30
        tasks[task_id]["stage"] = "Запуск 3D-генерации (Stage 1 & Stage 2)..."
        
        # 2. Forward pass
        outputs = model(
            processed_image,
            seed=int(params["seed"]),
            ss_steps=int(params["ss_steps"]),
            ss_cfg=float(params["ss_cfg"]),
            slat_steps=int(params["slat_steps"]),
            slat_cfg=float(params["slat_cfg"]),
            formats=['mesh', 'gaussian'],
            preprocess=False, # Already preprocessed
            refine_gs=params["refine_gs"],
            refine_steps=int(params["refine_steps"])
        )
        
        tasks[task_id]["progress"] = 70
        tasks[task_id]["stage"] = "Рендеринг превью 3D-гауссианов..."
        
        # 3. Render refined preview
        from trellis.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
        from trellis.renderers.gaussian_render import GaussianRenderer
        
        renderer = GaussianRenderer({
            "resolution": 512,
            "near": 0.8,
            "far": 1.6,
            "ssaa": 1,
            "bg_color": (0, 0, 0),
        })
        extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics([0.0], [0.0], 2.0, 40.0)
        
        with torch.no_grad():
            rendered = renderer.render(outputs['gaussian'], extrinsics[0], intrinsics[0])['color']
            
        rendered_np = (rendered.permute(1, 2, 0).cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
        rendered_img = Image.fromarray(rendered_np)
        
        refined_render_path = os.path.join(task_dir, "refined_render.png")
        rendered_img.save(refined_render_path)
        
        tasks[task_id]["progress"] = 85
        tasks[task_id]["stage"] = "Экспорт полигональной сетки и запекание текстур (GLB)..."
        
        # 4. Save aligned, textured mesh (GLB) + Gaussians (PLY) + metadata
        model.save_hybrid_asset(
            outputs,
            task_dir,
            prefix="forma_web",
            upscale_factor=int(params["upscale_factor"]),
            upscale_device=params["upscale_device"]
        )
        
        # Paths for result serving
        tasks[task_id]["progress"] = 100
        tasks[task_id]["stage"] = "Генерация успешно завершена!"
        tasks[task_id]["status"] = "completed"
        tasks[task_id]["result"] = {
            "glb_url": f"/static/outputs/{task_id}/forma_web.glb",
            "ply_url": f"/static/outputs/{task_id}/forma_web.ply",
            "obj_url": f"/static/outputs/{task_id}/forma_web_obj.zip",
            "meta_url": f"/static/outputs/{task_id}/forma_web_metadata.json",
            "mask_url": f"/static/outputs/{task_id}/preprocessed.png",
            "render_url": f"/static/outputs/{task_id}/refined_render.png"
        }
        print(f"[Backend] Task {task_id} completed successfully.")
        
    except Exception as e:
        import traceback
        err = traceback.format_exc()
        print(f"[Backend ERROR] Task {task_id} failed: {err}")
        tasks[task_id]["status"] = "failed"
        tasks[task_id]["stage"] = "Ошибка генерации"
        tasks[task_id]["error"] = str(e)

@app.post("/api/generate")
async def generate_3d(
    background_tasks: BackgroundTasks,
    image: UploadFile = File(...),
    seed: int = Form(42),
    ss_steps: int = Form(12),
    ss_cfg: float = Form(7.5),
    slat_steps: int = Form(12),
    slat_cfg: float = Form(3.0),
    refine_gs: bool = Form(False),
    refine_steps: int = Form(100),
    upscale_factor: int = Form(2),
    upscale_device: str = Form("Auto-Fallback")
):
    try:
        # Load and validate image
        contents = await image.read()
        import io
        img = Image.open(io.BytesIO(contents))
    except Exception as e:
        raise HTTPException(status_code=400, detail="Invalid image file.")
        
    task_id = str(uuid.uuid4())
    tasks[task_id] = {
        "status": "pending",
        "progress": 0,
        "stage": "Инициализация задачи в очереди...",
        "error": None,
        "result": None
    }
    
    params = {
        "seed": seed,
        "ss_steps": ss_steps,
        "ss_cfg": ss_cfg,
        "slat_steps": slat_steps,
        "slat_cfg": slat_cfg,
        "refine_gs": refine_gs,
        "refine_steps": refine_steps,
        "upscale_factor": upscale_factor,
        "upscale_device": upscale_device
    }
    
    # Launch pipeline run in background task thread
    background_tasks.add_task(run_pipeline_task, task_id, img, params)
    
    return {"task_id": task_id, "status": "pending"}

@app.get("/api/task/{task_id}")
async def get_task_status(task_id: str):
    if task_id not in tasks:
        raise HTTPException(status_code=404, detail="Task not found.")
    return tasks[task_id]

# Root endpoint redirect to index.html
@app.get("/")
async def read_index():
    return FileResponse(os.path.join(BASE_DIR, "static/index.html"))

# Mount static folder
app.mount("/static", StaticFiles(directory="static"), name="static")

if __name__ == "__main__":
    import uvicorn
    # Start FastAPI server on port 7860
    uvicorn.run(app, host="127.0.0.1", port=7860)