File size: 4,029 Bytes
04de76b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
"""FastAPI backend for InteriorFusion inference service."""

import os
import tempfile
from pathlib import Path
from typing import List, Optional

import torch
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from PIL import Image
import io
import base64

from interiorfusion.pipelines import InteriorFusionPipeline, InteriorFusionOutput

app = FastAPI(
    title="InteriorFusion API",
    description="Single image to 3D interior scene generation",
    version="0.1.0",
)

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

# Global pipeline instance
_pipeline: Optional[InteriorFusionPipeline] = None


def get_pipeline() -> InteriorFusionPipeline:
    global _pipeline
    if _pipeline is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        _pipeline = InteriorFusionPipeline(
            model_size="L",
            device=device,
            dtype=torch.float16,
        )
    return _pipeline


@app.post("/generate")
async def generate_3d_scene(
    image: UploadFile = File(...),
    room_type: Optional[str] = Form(None),
    style: Optional[str] = Form(None),
    formats: str = Form("glb,ply"),
    model_size: str = Form("L"),
):
    """
    Generate a 3D interior scene from a single image.
    
    Returns download links for the generated 3D files.
    """
    # Read image
    contents = await image.read()
    img = Image.open(io.BytesIO(contents)).convert("RGB")
    
    # Parse formats
    output_formats = [f.strip() for f in formats.split(",")]
    
    # Run pipeline
    pipeline = get_pipeline()
    output = pipeline(
        image=img,
        room_type_hint=room_type,
        style_hint=style,
    )
    
    # Export
    output_dir = tempfile.mkdtemp()
    output.export_all(output_dir)
    
    # Collect file paths
    files = {}
    for fmt in output_formats:
        path = Path(output_dir) / f"scene.{fmt}"
        if path.exists():
            files[fmt] = str(path)
    
    return JSONResponse({
        "success": True,
        "room_type": output.room_type,
        "style": output.style,
        "processing_time": output.processing_time,
        "num_objects": len(output.object_meshes),
        "files": files,
    })


@app.post("/edit")
async def edit_scene(
    scene_glb: UploadFile = File(...),
    edit_action: str = Form(...),  # "move", "replace", "remove", "add"
    object_id: Optional[int] = Form(None),
    new_image: Optional[UploadFile] = File(None),
    position: Optional[str] = Form(None),  # JSON array [x, y, z]
):
    """
    Edit an existing scene.
    
    Actions:
    - move: Move an existing object
    - replace: Replace an object with a new one
    - remove: Remove an object
    - add: Add a new object
    """
    import json
    
    pipeline = get_pipeline()
    
    # Parse position
    pos = None
    if position:
        pos = json.loads(position)
    
    # Build edit dict
    edit = {"action": edit_action}
    if object_id is not None:
        edit["object_id"] = object_id
    if pos:
        edit["position"] = pos
    if new_image:
        contents = await new_image.read()
        edit["new_image"] = Image.open(io.BytesIO(contents)).convert("RGB")
    
    # For simplicity, return not-implemented
    return JSONResponse({
        "success": False,
        "message": "Scene editing API coming soon",
    })


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "ok", "version": "0.1.0"}


@app.get("/download/{filename}")
async def download_file(filename: str):
    """Download a generated file."""
    # In production, use proper file storage (S3, etc.)
    # For now, placeholder
    return JSONResponse({"message": f"Download {filename} from storage"})


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