Spaces:
Running
Running
File size: 2,678 Bytes
f3b96de |
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 |
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import FileResponse
import os
import shutil
import uuid
import requests
from typing import Optional
from app.utils import run_inference
from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN")
print(hf_token)
login(token=hf_token)
app = FastAPI(title="Stable Fast 3D API")
@app.get("/")
async def root():
return {"message": "Welcome to Stable Fast 3D API. Use /generate-3d endpoint to convert 2D images to 3D models."}
@app.post("/generate-3d/")
async def generate_3d_model_upload(image: UploadFile = File(...)):
"""Generate 3D model from uploaded image file"""
return await process_image(image=image)
@app.get("/generate-3d/")
async def generate_3d_model_url(image_url: str = Query(..., description="URL of the image to convert to 3D")):
"""Generate 3D model from image URL"""
return await process_image(image_url=image_url)
async def process_image(image: Optional[UploadFile] = None, image_url: Optional[str] = None):
# Create unique ID for this request
temp_id = str(uuid.uuid4())
input_path = f"/app/tmp/{temp_id}.png"
output_dir = f"/app/tmp/{temp_id}_output"
os.makedirs(output_dir, exist_ok=True)
try:
# Handle image from upload or URL
if image:
with open(input_path, "wb") as f:
shutil.copyfileobj(image.file, f)
elif image_url:
response = requests.get(image_url, stream=True)
if response.status_code != 200:
raise HTTPException(status_code=400, detail="Could not download image from URL")
with open(input_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
else:
raise HTTPException(status_code=400, detail="Either image file or image URL must be provided")
# Run the inference
glb_path = run_inference(input_path, output_dir)
if not os.path.exists(glb_path):
raise HTTPException(status_code=500, detail="Failed to generate 3D model")
# Return the GLB file
return FileResponse(
path=glb_path,
media_type="model/gltf-binary",
filename="model.glb",
headers={"Content-Disposition": f"attachment; filename=model.glb"}
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
finally:
# Clean up temporary files
if os.path.exists(input_path):
os.remove(input_path) |