Feat/gateway (#7)

* feat: gateway blocking pipeline

* feat: spz auto update

* Implement SPZLoader

* feat: unblocking operator

* fix: correct name for parallel generations

* refactor: remove unnecessary print

* fix: version fix

* fix import issue in dependencies.py

---------

Co-authored-by: Denis Avvakumov <denisavvakumov@gmail.com>
Co-authored-by: Ruben Hohndorf <ruben.hohndorf@gmail.com>

Author: 3 people

.gitignore

@@ -2,6 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
+.env
 
 # C extensions
 *.so
@@ -131,4 +132,5 @@ dmypy.json
 # Vim stuff
 *.sw*
 
-*.blend1
+*.blend1
+spz_version.txt

fourofour_3d_gen/__init__.py

@@ -7,6 +7,7 @@
 def register():
     preferences.register()
     modules.append(preferences)
+    
     if dependencies.installed():
         from . import props, ui, ops
 
@@ -17,6 +18,17 @@ def register():
         modules.append(ui)
         modules.append(props)
 
+        try:
+            from .spz_loader import init_spz
+            from pathlib import Path
+
+            pkg_dir = Path(__file__).resolve().parent
+            print(f"Initializing SPZ with library path: {pkg_dir}")
+            init_spz(str(pkg_dir))
+        except Exception as e:
+            # Report failure to load native SPZ library to console only
+            print(f"SPZ initialization failed: {e}")
+
 
 def unregister():
     for m in modules:

fourofour_3d_gen/blender_manifest.toml

@@ -1,7 +1,7 @@
 schema_version = "1.0.0"
 
 id = "fourofour_3d_gen"
-version = "0.9.0"
+version = "0.10.0"
 name = "404 3D Generator"
 tagline = "404 3D Generator Extension"
 maintainer = "Ruben Hohndorf <ruben@atlas.design>"

fourofour_3d_gen/client.py

@@ -2,53 +2,80 @@
 import base64
 import json
 import tempfile
-import websocket
-
+import time
+from .spz_loader import get_spz
 
 from .protocol import Auth, PromptData, TaskStatus, TaskUpdate
+from .gateway.gateway_api import GatewayApi
+from .gateway.gateway_task import GatewayTask, GatewayTaskStatusResponse, GatewayTaskStatus
+
+
+_GATEWAY_STATUS_CHECK_INTERVAL_SEC: int = 5
+_GATEWAY_TASK_TIMEOUT_SEC: int = 10 * 60 # 10 minutes
+
+
+class GatewayErrorBase(Exception):
+    pass
+
+
+class GatewayTimeoutError(GatewayErrorBase):
+    pass
+
+
+class GatewayFailureError(GatewayErrorBase):
+    pass
+
+
+class Client:
+
+    def __init__(self) -> None:
+        self._task: GatewayTask | None = None
+        self._task_start_time: float | None = None
+        gateway_url = bpy.context.preferences.addons[__package__].preferences.url
+        gateway_api_key = bpy.context.preferences.addons[__package__].preferences.token
+        self._gateway_api: GatewayApi = GatewayApi(gateway_url=gateway_url, gateway_api_key=gateway_api_key)
+
+    @property
+    def task_id(self) -> str | None:
+        return self._task.id if self._task else None
+    
+    @property
+    def prompt(self) -> str | None:
+        return self._task.prompt if self._task else None
+
+    def request_model(self, prompt: str) -> None:
+        task = self._gateway_api.add_task(text_prompt=prompt)
+        print(f"Task added: {task.id}")
+        self._task = task
+        self._task_start_time = time.time()
+    
+    def get_result(self) -> str | None:
+        if self._task is None or self._task_start_time is None:
+            return None
+            
+        # Check for timeout
+        if time.time() - self._task_start_time > _GATEWAY_TASK_TIMEOUT_SEC:
+            raise GatewayTimeoutError("Gateway timeout error")
+
+        task_status_response = self._gateway_api.get_status(task=self._task)
+        task_status = task_status_response.status
+
+        if task_status not in [GatewayTaskStatus.SUCCESS, GatewayTaskStatus.FAILURE]:
+            return None
+        
+        if task_status == GatewayTaskStatus.FAILURE:
+            raise GatewayFailureError(f"Gateway failure error")
+
+        if task_status == GatewayTaskStatus.SUCCESS:
+            spz_data = self._gateway_api.get_result(task=self._task)
+            print(f"Received result for task: {self._task.id}")
 
+            loader = get_spz()
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".ply") as temp_file:
+                ply_data = loader.decompress(spz_data, include_normals=False)
+                temp_file.write(ply_data)
+                filepath = temp_file.name
+                print(f"Saved result to: {filepath}")
+            return filepath
 
-def request_model(prompt: str) -> tuple[None, None] | tuple[str, str]:
-    url = bpy.context.preferences.addons[__package__].preferences.url
-    api_key = bpy.context.preferences.addons[__package__].preferences.token
-    filepath = None
-    winner_hotkey = None
-
-    def on_message(ws, message):
-        nonlocal filepath, winner_hotkey
-        update = TaskUpdate(**json.loads(message))
-        if update.status == TaskStatus.STARTED:
-            print("Task started")
-        elif update.status == TaskStatus.FIRST_RESULTS:
-            score = update.results.score if update.results else None
-            assets = update.results.assets or "" if update.results else ""
-            print(f"First results. Score: {score}. Size: {len(assets)}")
-        elif update.status == TaskStatus.BEST_RESULTS:
-            score = update.results.score if update.results else None
-            assets = update.results.assets or "" if update.results else ""
-            print(f"Best results. Score: {score}. Size: {len(assets)}")
-            print(f"Stats: {update.statistics}")
-
-            if assets:
-                with tempfile.NamedTemporaryFile(delete=False, suffix=".ply") as temp_file:
-                    temp_file.write(base64.b64decode(assets.encode("utf-8")))
-                    filepath = temp_file.name
-                    ws.close()
-
-    def on_error(ws, error):
-        print(f"WebSocket connection error: {error}")
-
-    def on_close(ws, close_status_code, close_msg):
-        print(f"WebSocket connection closed: {close_status_code} {close_msg}")
-
-    def on_open(ws):
-        auth_data = Auth(api_key=api_key).dict()
-        prompt_data = PromptData(prompt=prompt, send_first_results=True).dict()
-        ws.send(json.dumps(auth_data))
-        ws.send(json.dumps(prompt_data))
-
-    ws = websocket.WebSocketApp(url, on_message=on_message, on_error=on_error, on_close=on_close)
-    ws.on_open = on_open
-    ws.run_forever()
-
-    return (filepath, winner_hotkey)
+        return None

fourofour_3d_gen/dependencies.py

@@ -8,7 +8,6 @@ def installed() -> bool:
     try:
         import numpy
         import pydantic
-        import websocket
         import mixpanel
 
         return True
@@ -44,15 +43,20 @@ def install_dependencies_from_requirements():
         env=env_var,
     )
 
+def update_spz():
+    from .spz_updater import SPZUpdater
+
+    if SPZUpdater.need_update():
+        SPZUpdater.update()
+
 
 def install() -> None:
     install_pip()
-
     install_dependencies_from_requirements()
 
-
 if __name__ == "__main__":
     if installed():
         print("dependencies installed")
     else:
         install()
+    update_spz()

fourofour_3d_gen/gateway/gateway_api.py

@@ -0,0 +1,84 @@
+import json
+import requests
+
+from typing import Any, cast
+from urllib.parse import urlencode
+
+from .gateway_routes import GatewayRoutes
+from .gateway_task import GatewayTask, GatewayTaskStatusResponse
+
+
+class GatewayErrorBase(Exception):
+    pass
+
+
+class GatewayAddTaskError(GatewayErrorBase):
+    pass
+
+
+class GatewayGetStatusError(GatewayErrorBase):
+    pass
+
+
+class GatewayGetResultError(GatewayErrorBase):
+    pass
+
+
+class GatewayNoAttachmentError(GatewayErrorBase):
+    pass
+
+
+class GatewayApi:
+    """API client for interacting with gateway."""
+
+    def __init__(self, *, gateway_url: str, gateway_api_key: str) -> None:
+        self._http2_client = requests.Session()
+        """HTTP3 client for interacting with gateway."""
+        self._gateway_url = gateway_url
+        """URL of the gateway."""
+        self._gateway_api_key = gateway_api_key
+        """API key of the gateway that provides no rate limit for discord bot."""
+
+    def add_task(self, *, text_prompt: str) -> GatewayTask:
+        """Adds a task to the gateway."""
+        try:
+            url = self._construct_url(host=self._gateway_url, route=GatewayRoutes.ADD_TASK)
+            payload = {"prompt": text_prompt}
+            headers = {"x-api-key": self._gateway_api_key}
+            response = self._http2_client.post(url=url, json=payload, headers=headers)
+            print(response.text)
+            response.raise_for_status()
+            data = json.loads(response.text)
+            return GatewayTask.model_validate(data)
+        except Exception as e:
+            raise GatewayAddTaskError(f"Gateway: error to add task: {e}") from e
+
+    def get_status(self, *, task: GatewayTask) -> GatewayTaskStatusResponse:
+        """Gets the status of a task."""
+        try:
+            url = self._construct_url(host=self._gateway_url, route=GatewayRoutes.GET_STATUS, id=task.id)
+            headers = {"x-api-key": self._gateway_api_key}
+            response = self._http2_client.get(url=url, headers=headers)
+            response.raise_for_status()
+            data = json.loads(response.text)
+            reason = data.get("reason", None)
+            return GatewayTaskStatusResponse(status=data["status"], reason=reason)
+        except Exception as e:
+            raise GatewayGetStatusError(f"Gateway: error to get status: {e}") from e
+
+    def get_result(self, *, task: GatewayTask) -> bytes:
+        """Gets generated 3D asset in spz format."""
+        try:
+            url = self._construct_url(host=self._gateway_url, route=GatewayRoutes.GET_RESULT, id=task.id)
+            headers = {"x-api-key": self._gateway_api_key}
+            response = self._http2_client.get(url=url, headers=headers)
+            response.raise_for_status()
+            if response.headers.get('content-disposition', '').startswith('attachment'):
+                return cast(bytes, response.content)
+            raise GatewayNoAttachmentError()
+        except Exception as e:
+            raise GatewayGetResultError(f"Gateway: error to get result: {e}") from e
+
+    def _construct_url(self, *, host: str, route: GatewayRoutes, **kwargs: Any) -> str:
+        return f"{host}{route.value}?{urlencode(kwargs)}"
+

fourofour_3d_gen/gateway/gateway_routes.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+
+class GatewayRoutes(Enum):
+    """Routes for gateway nodes"""
+
+    ADD_TASK = "/add_task"
+    """Send text prompt to gateway to generate 3D asset."""
+    GET_STATUS = "/get_status"
+    """Get status of the task."""
+    GET_RESULT = "/get_result"
+    """Get result of the generation in spz format."""

fourofour_3d_gen/gateway/gateway_task.py

@@ -0,0 +1,35 @@
+from enum import Enum
+
+from pydantic import BaseModel
+
+
+class GatewayTaskStatus(Enum):
+    """Status of the task in gateway"""
+
+    NO_RESULT = "NoResult"
+    FAILURE = "Failure"
+    PARTIAL_RESULT = "PartialResult"
+    SUCCESS = "Success"
+
+
+class GatewayTaskStatusResponse(BaseModel):
+    """Response from the gateway for the task status"""
+
+    status: GatewayTaskStatus
+    """Status of the task in gateway"""
+    reason: str | None = None
+    """Reason for the status. If status is not SUCCESS, reason is not None."""
+
+
+class GatewayTask(BaseModel):
+    """Task for the gateway"""
+
+    id: str
+    """Unique id of the task in gateway"""
+    prompt: str
+    """Text prompt to generate the asset 3D asset"""
+    # TODO: in what format should we store in bunny net.
+    result: bytes | None = None
+    """Result of the task in gateway in spz format"""
+    task_status: GatewayTaskStatus = GatewayTaskStatus.NO_RESULT
+    """Status of the task in gateway"""

fourofour_3d_gen/gaussian_splatting.py

@@ -9,7 +9,7 @@
 RECOMMENDED_MAX_GAUSSIANS = 200_000
 
 
-def import_gs(filepath: str, name: str, winner_hotkey: str = ""):
+def import_gs(filepath: str, name: str):
 
     if "GaussianSplatting" not in bpy.data.node_groups:
         script_file = os.path.realpath(__file__)
@@ -131,8 +131,6 @@ def import_gs(filepath: str, name: str, winner_hotkey: str = ""):
     obj.rotation_euler = (-np.pi / 2, 0, 0)
     obj.rotation_euler[0] = 1.5708
 
-    obj["Bittensor Miner"] = winner_hotkey
-
     print("Mesh attributes added in", time.time() - start_time, "seconds")
 
     setup_nodes(obj)