{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6117abdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, requests, time\n",
    "\n",
    "JOB_ID    = os.environ[\"JOB_ID\"]\n",
    "SPACE_URL = os.environ[\"SPACE_URL\"].rstrip(\"/\")\n",
    "\n",
    "print(f\"Job: {JOB_ID}\")\n",
    "print(f\"Space: {SPACE_URL}\")\n",
    "\n",
    "# 1. Download image\n",
    "print(\"Downloading image...\")\n",
    "r = requests.get(f\"{SPACE_URL}/image/{JOB_ID}\", timeout=30)\n",
    "r.raise_for_status()\n",
    "with open(\"/kaggle/working/input.png\", \"wb\") as f:\n",
    "    f.write(r.content)\n",
    "print(\"Image downloaded!\")\n",
    "\n",
    "# 2. Remove background\n",
    "os.system(\"pip install -q transparent-background gradio_client pillow\")\n",
    "from transparent_background import Remover\n",
    "from PIL import Image\n",
    "\n",
    "print(\"Removing background...\")\n",
    "remover = Remover()\n",
    "img = Image.open(\"/kaggle/working/input.png\").convert(\"RGB\")\n",
    "out = remover.process(img, type=\"rgba\")\n",
    "out.save(\"/kaggle/working/input_nobg.png\")\n",
    "print(\"BG removed!\")\n",
    "\n",
    "# 3. Run TRELLIS\n",
    "from gradio_client import Client, handle_file\n",
    "\n",
    "MAX_RETRIES = 3\n",
    "result = None\n",
    "\n",
    "for attempt in range(1, MAX_RETRIES + 1):\n",
    "    try:\n",
    "        print(f\"Connecting to TRELLIS (attempt {attempt}/{MAX_RETRIES})...\")\n",
    "        client = Client(\"trellis-community/TRELLIS\")\n",
    "        client.predict(api_name=\"/start_session\")\n",
    "        print(\"Session ready! Generating...\")\n",
    "\n",
    "        result = client.predict(\n",
    "            image=handle_file(\"/kaggle/working/input_nobg.png\"),\n",
    "            multiimages=[],\n",
    "            seed=0,\n",
    "            ss_guidance_strength=7.5,\n",
    "            ss_sampling_steps=12,\n",
    "            slat_guidance_strength=3.0,\n",
    "            slat_sampling_steps=12,\n",
    "            multiimage_algo=\"stochastic\",\n",
    "            mesh_simplify=0.95,\n",
    "            texture_size=1024,\n",
    "            api_name=\"/generate_and_extract_glb\"\n",
    "        )\n",
    "        print(\"Generation done!\")\n",
    "        break\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Attempt {attempt} failed: {e}\")\n",
    "        if attempt < MAX_RETRIES:\n",
    "            time.sleep(30)\n",
    "        else:\n",
    "            raise RuntimeError(f\"TRELLIS failed after {MAX_RETRIES} attempts: {e}\")\n",
    "\n",
    "# 4. POST GLB back\n",
    "glb_src = result[1] or result[2]\n",
    "print(f\"Sending GLB ({os.path.getsize(glb_src)/1024/1024:.1f} MB)...\")\n",
    "\n",
    "with open(glb_src, \"rb\") as f:\n",
    "    resp = requests.post(\n",
    "        f\"{SPACE_URL}/receive_glb\",\n",
    "        params={\"job_id\": JOB_ID},\n",
    "        files={\"file\": (\"mesh.glb\", f, \"model/gltf-binary\")},\n",
    "        timeout=120\n",
    "    )\n",
    "\n",
    "resp.raise_for_status()\n",
    "print(f\"GLB delivered! Job {JOB_ID} complete.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}