Create app.py

app.py

@@ -0,0 +1,165 @@
+import io
+import zipfile
+from typing import Dict
+
+import torch
+from fastapi import FastAPI, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
+from PIL import Image
+from rembg import remove
+from diffusers import DiffusionPipeline
+
+
+app = FastAPI(title="Zero123++ Inference API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+MODEL_ID = "sudo-ai/zero123plus-v1.2"
+CUSTOM_PIPELINE = "sudo-ai/zero123plus-pipeline"
+
+pipeline = None
+
+
+def load_pipeline():
+    global pipeline
+
+    if pipeline is not None:
+        return pipeline
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "CUDA GPU is not available. Please enable GPU hardware on the Hugging Face Space."
+        )
+
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        custom_pipeline=CUSTOM_PIPELINE,
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+    )
+
+    pipe.to("cuda")
+    pipe.enable_attention_slicing()
+
+    pipeline = pipe
+    return pipeline
+
+
+def image_to_bytes(image: Image.Image, fmt: str = "PNG") -> bytes:
+    buffer = io.BytesIO()
+    image.save(buffer, format=fmt)
+    buffer.seek(0)
+    return buffer.getvalue()
+
+
+def crop_selected_views(grid: Image.Image) -> Dict[str, Image.Image]:
+    """
+    Zero123++ output is expected as a 2-column x 3-row grid.
+    We keep views 1, 3, 5, and 6 for LGM.
+    """
+
+    grid = grid.convert("RGB")
+    w, h = grid.size
+
+    cols, rows = 2, 3
+    tile_w, tile_h = w // cols, h // rows
+
+    selected = {
+        1: "front_right",
+        3: "right",
+        5: "back_right",
+        6: "back_left",
+    }
+
+    outputs = {}
+
+    for idx_1based, name in selected.items():
+        row = (idx_1based - 1) // cols
+        col = (idx_1based - 1) % cols
+
+        box = (
+            col * tile_w,
+            row * tile_h,
+            (col + 1) * tile_w,
+            (row + 1) * tile_h,
+        )
+
+        tile = grid.crop(box).resize((256, 256), Image.LANCZOS)
+
+        # Remove background and paste on white
+        tile_rgba = remove(tile.convert("RGBA"))
+        white_bg = Image.new("RGBA", tile_rgba.size, (255, 255, 255, 255))
+        white_bg.paste(tile_rgba, mask=tile_rgba.split()[3])
+
+        outputs[f"view_{idx_1based}_{name}.png"] = white_bg.convert("RGB")
+
+    return outputs
+
+
+@app.get("/")
+def root():
+    return {
+        "status": "running",
+        "service": "zero123plus-inference",
+        "model": MODEL_ID,
+        "output": "6-view grid + cropped views 1, 3, 5, 6 for LGM",
+    }
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_device": torch.cuda.get_device_name(0)
+        if torch.cuda.is_available()
+        else None,
+    }
+
+
+@app.post("/generate")
+async def generate(file: UploadFile = File(...), steps: int = 75):
+    try:
+        pipe = load_pipeline()
+
+        contents = await file.read()
+        input_image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        with torch.inference_mode():
+            result = pipe(input_image, num_inference_steps=steps).images[0]
+
+        cropped_views = crop_selected_views(result)
+
+        zip_buffer = io.BytesIO()
+
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+            zf.writestr("multiview_grid.png", image_to_bytes(result))
+
+            for filename, image in cropped_views.items():
+                zf.writestr(filename, image_to_bytes(image))
+
+        zip_buffer.seek(0)
+
+        return StreamingResponse(
+            zip_buffer,
+            media_type="application/zip",
+            headers={
+                "Content-Disposition": "attachment; filename=zero123plus_outputs.zip"
+            },
+        )
+
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "error": str(e),
+                "message": "Zero123++ generation failed. Check Space logs for details.",
+            },
+        )