webgpu 0.0.1.dev0__py3-none-any.whl

webgpu/link/link.js

@@ -0,0 +1,431 @@
+/* eslint-disable */
+
+function serializeEvent(event) {
+  event.preventDefault();
+  const keys = [
+    'button',
+    'altKey',
+    'metaKey',
+    'ctrlKey',
+    'shiftKey',
+    'x',
+    'y',
+    'deltaX',
+    'deltaY',
+    'deltaMode',
+    'movementX',
+    'movementY',
+  ];
+  return Object.fromEntries(keys.map((k) => [k, event[k]]));
+}
+
+function decodeB64(data) {
+  return Uint8Array.from(atob(data), (c) => c.charCodeAt(0)).buffer;
+}
+
+function encodeB64(buffer) {
+  return btoa(String.fromCharCode.apply(null, new Uint8Array(buffer)));
+}
+
+function isPrimitive(value) {
+  return (
+    value === null || (typeof value !== 'object' && typeof value !== 'function')
+  );
+}
+
+function isPrimitiveDict(obj) {
+  if (obj.$children) return false;
+  if (obj === window) return false;
+  if (obj === navigator) return false;
+
+  if (typeof obj !== 'object' || obj === null) return false;
+  if (Array.isArray(obj)) {
+    return obj.every(isPrimitiveDict);
+  }
+
+  for (let key in obj) {
+    if (Object.prototype.hasOwnProperty.call(obj, key)) {
+      const value = obj[key];
+
+      if (!isPrimitive(value) && !isPrimitiveDict(value)) {
+        return false;
+      }
+    }
+  }
+
+  if (obj.constructor !== Object) return false;
+
+  return true;
+}
+
+function createProxy(link, id, parent_id) {
+  const target = function (...args) {
+    return link.call(id, args, parent_id);
+  };
+  target.handleEvent = async (eventType, event) => {
+    return link.call(id, [eventType, event], parent_id, '_handle');
+  };
+  target.callFunction = (self, arg) => {
+    return link.call(id, [arg], parent_id);
+  };
+  target.callMethod = async (prop, arg) => {
+    return await link.callMethod(id, prop, [arg]);
+  };
+  target.callMethodIgnoreResult = (prop, arg) => {
+    link.callMethodIgnoreResult(id, prop, [arg]);
+  };
+  const handler = {
+    get: function (obj, prop) {
+      if (prop === '_parent_id') return parent_id;
+      if (prop === '_id') return id;
+      if (prop === '_link') return link;
+      if (prop === 'handleEvent') return target.handleEvent;
+      if (prop === 'call') {
+        return target.callFunction;
+      }
+      if (prop === 'then') return Reflect.get(...arguments);
+      if (prop.startsWith('__')) return Reflect.get(...arguments);
+      if (prop === 'callMethodIgnoreResult') return Reflect.get(...arguments);
+      if (prop === 'callMethod') return Reflect.get(...arguments);
+
+      return link.getProp(id, prop);
+    },
+    set: function (obj, prop, value) {
+      link.setProp(id, prop, value);
+    },
+
+    apply: function (obj, thisArg, args) {
+      return link.call(id, args, parent_id);
+    },
+  };
+  return new Proxy(target, handler);
+}
+
+class CrossLink {
+  constructor(connection) {
+    this.requestCounter = 1;
+    this.requests = {};
+
+    this.counter = 1;
+    this.objects = {};
+
+    this.connection = connection;
+    this.connection.onMessage((data) => this.onMessage(data));
+    this.connected = new Promise((resolve) => {
+      this.connection.onOpen(() => {
+        console.log('connection open');
+        resolve();
+      });
+    });
+  }
+
+  async _sendRequestAwaitResponse(data) {
+    //console.log('sending request', data);
+    const request_id = this.requestCounter++;
+    data.request_id = request_id;
+    try {
+      const result = await new Promise((resolve, reject) => {
+        this.requests[request_id] = resolve;
+        this.connection.send(JSON.stringify(data));
+        setTimeout(() => {
+          reject(
+            new Error(
+              `Timeout, request ${request_id}, data: ${JSON.stringify(data)}`
+            )
+          );
+        }, 10000);
+      });
+      // const t = Date.now() - requestData.sent;
+      return result;
+    } finally {
+      delete this.requests[request_id];
+    }
+  }
+
+  async getProp(id, prop) {
+    return await this._sendRequestAwaitResponse({ type: 'get', id, prop });
+  }
+
+  async getItem(id, key) {
+    return await this._sendRequestAwaitResponse({ type: 'get', id, key });
+  }
+
+  async setProp(id, prop, value) {
+    this.connection.send({
+      type: 'set',
+      id,
+      prop,
+      value: await this._dumpData(value),
+    });
+  }
+
+  async setItem(id, key, value) {
+    this.connection.send({
+      type: 'set',
+      id,
+      key,
+      value: await this._dumpData(value),
+    });
+  }
+
+  async call(id, args = [], parent_id = undefined, prop = undefined) {
+    return await this._sendRequestAwaitResponse({
+      type: 'call',
+      id,
+      parent_id,
+      args: await this._dumpData(args),
+      prop,
+    });
+  }
+
+  async callMethod(id, prop, args = []) {
+    return await this._sendRequestAwaitResponse({
+      type: 'call',
+      id,
+      args: await this._dumpData(args),
+      prop,
+    });
+  }
+
+  async callMethodIgnoreResult(id, prop, args = []) {
+    return await this.connection.send(
+      JSON.stringify({
+        type: 'call',
+        id,
+        args: await this._dumpData(args),
+        prop,
+      })
+    );
+  }
+
+  expose(name, obj) {
+    this.objects[name] = obj;
+  }
+
+  async _dumpData(data) {
+    //console.log("dumping data", data);
+    if (data === null) return undefined;
+
+    if (isPrimitive(data)) return data;
+
+    if (data instanceof MouseEvent) return serializeEvent(data);
+    if (data instanceof Event) return serializeEvent(data);
+    if (data instanceof InputEvent) return serializeEvent(data);
+
+    if (data instanceof ArrayBuffer)
+      return {
+        __is_crosslink_type__: true,
+        type: 'bytes',
+        value: encodeB64(data),
+      };
+
+    if (data.constructor === Array) {
+      const result = [];
+      for (let item of data) result.push(await this._dumpData(item));
+      return result;
+    }
+
+    if (data.__is_crosslink_type__) return data;
+
+    if (isPrimitiveDict(data)) {
+      return data;
+    }
+    if (data.constructor === Object) {
+      const result = {};
+      Object.keys(data).map(async (key) => {
+        result[key] = await this._dumpData(data[key]);
+      });
+      return result;
+    }
+
+    // complex type - store it in objects only send its id
+    const id = this.counter++;
+    this.objects[id] = data;
+    return {
+      __is_crosslink_type__: true,
+      type: 'proxy',
+      id,
+    };
+  }
+
+  _loadValue(value) {
+    if (value === null || value === undefined) return undefined;
+    if (!value.__is_crosslink_type__) return value;
+
+    if (value.type == 'bytes') return decodeB64(value.value);
+    if (value.type == 'object') return this.objects[value.id];
+    if (value.type == 'proxy')
+      return createProxy(this, value.id, value.parent_id);
+
+    console.error('Cannot load value, unknown value type:', value);
+  }
+
+  _loadData(data) {
+    if (
+      data === undefined ||
+      data === null ||
+      typeof data !== 'object' ||
+      data.__is_crosslink_type__
+    )
+      return this._loadValue(data);
+
+    Object.keys(data).map((key) => {
+      data[key] = this._loadData(data[key]);
+    });
+    return data;
+  }
+
+  async sendResponse(data, request_id, parent_id) {
+    if (request_id === undefined) {
+      return;
+    }
+
+    const value = await this._dumpData(await Promise.resolve(data));
+    //console.log('encoded data', value);
+
+    this.connection.send(
+      JSON.stringify({
+        type: 'response',
+        request_id,
+        value,
+      })
+    );
+  }
+
+  async onMessage(event) {
+    const data = JSON.parse(event.data);
+    const request_id = data.request_id;
+
+    let obj = data.id ? this.objects[data.id] : window;
+    // console.log('onMessage', data, obj);
+
+    let response = null;
+
+    switch (data.type) {
+      case 'call':
+        const args = this._loadData(data.args);
+        let self = null;
+        if (data.prop) {
+          self = this.objects[data.id];
+          obj = self[data.prop];
+        } else {
+          self = this.objects[data.parent_id];
+        }
+
+        // console.log('calling', 'data', data, 'obj', obj, 'self', self, 'args', args);
+
+        response = obj.apply(self, args);
+        break;
+      case 'get_keys':
+        response = Object.keys(obj);
+        break;
+      case 'get':
+        if (data.prop) response = obj[data.prop];
+        else if (data.key) response = obj[data.key];
+        else response = obj;
+        if (response && (data.prop || data.key)) {
+          response = await this._dumpData(response);
+          if (typeof response === 'object') response.parent_id = data.id;
+        }
+        break;
+      case 'set':
+        const value = this._loadData(data.value);
+        if (data.prop) obj[data.prop] = value;
+        if (data.key) obj[data.key] = value;
+        break;
+      case 'delete':
+        this.objects[data.id] = undefined;
+        break;
+      case 'response':
+        this.requests[request_id](this._loadData(data.value));
+        return;
+      default:
+        console.error('Unknown message type:', data, data.type);
+    }
+
+    if (request_id !== undefined && data.type !== 'response') {
+      response = await response;
+      this.sendResponse(response, request_id);
+    }
+  }
+}
+
+export function WebsocketLink(url) {
+  const socket = new WebSocket(url);
+  return new CrossLink({
+    send: (data) => socket.send(data),
+    onMessage: (callback) => (socket.onmessage = callback),
+    onOpen: (callback) => (socket.onopen = callback),
+  });
+}
+
+export function WebworkerLink(worker) {
+  // wait for first message, which means the worker has loaded and is ready
+  const workerReady = new Promise((resolve) => {
+    worker.addEventListener(
+      'message',
+      (event) => {
+        resolve();
+      },
+      { once: true }
+    );
+  });
+  return new CrossLink({
+    send: (data) => worker.postMessage(data),
+    onMessage: async (callback) => {
+      await workerReady;
+      worker.addEventListener('message', callback);
+    },
+    onOpen: async (callback) => {
+      await workerReady;
+      callback();
+    },
+  });
+}
+
+window.createLilGUI = async (args) => {
+  if (window.lil === undefined) {
+    const url = 'https://cdn.jsdelivr.net/npm/lil-gui@0.20';
+    if (window.define === undefined) {
+      await import(url);
+    } else {
+      await new Promise(async (resolve) => {
+        require([url], (module) => {
+          window.lil = module;
+          resolve();
+        });
+      });
+    }
+  }
+  return new window.lil.GUI(args);
+};
+
+window.importPackage=async (url) => {
+  if (window.define === undefined) {
+    await import(url);
+  } else {
+    await new Promise(async (resolve) => {
+      require([url], (module) => {
+        resolve(module);
+      });
+    });
+  }
+}
+
+window.patchedRequestAnimationFrame = (device, context, target) => {
+  // context.getCurrentTexture() is only guaranteed to be valid during the requestAnimationFrame callback
+  // Thus, in order to render from python asynchroniously, we are always rendering into a separate texture
+  // The actual callback here only copies the rendered image from the separate render target texture to the current texture
+  requestAnimationFrame((t) => {
+    const current = context.getCurrentTexture();
+    const encoder = device.createCommandEncoder();
+
+    encoder.copyTextureToTexture(
+      { texture: target },
+      { texture: current },
+      { width: current.width, height: current.height }
+    );
+
+    device.queue.submit([encoder.finish()]);
+  });
+};

webgpu/link/proxy.py

@@ -0,0 +1,81 @@
+from .base import LinkBase
+
+
+class _ProxyIterator:
+    def __init__(self, proxy):
+        self._proxy = proxy
+        self._keys = proxy._get_keys()
+        self._index = 0
+
+    def __next__(self):
+        if self._index < len(self._keys):
+            key = self._keys[self._index]
+            self._index += 1
+            return key
+        else:
+            raise StopIteration
+
+    def __iter__(self):
+        return self
+
+
+class Proxy:
+    _link: LinkBase
+    _parent_id: int | None
+    _id: int | None
+
+    def __init__(self, link, parent_id=None, id=None):
+        self._link = link
+        self._id = id
+        self._parent_id = parent_id
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
+
+    def __getattr__(self, key):
+        if (
+            key
+            in [
+                "_id",
+                "_parent_id",
+                "_link",
+                "_call",
+                "_call_method_ignore_return",
+                "_call_method",
+            ]
+            or isinstance(key, str)
+            and key.startswith("__")
+        ):
+            return super().__getattr__(key)
+        return self._link.get(self._id, key)
+
+    def __setattr__(self, key, value):
+        if key in ["_id", "_parent_id", "_link"]:
+            return super().__setattr__(key, value)
+
+        return self._link.set(self._id, key, value)
+
+    def __call__(self, *args, _ignore_result=False):
+        return self._link.call(self._id, list(args), self._parent_id, _ignore_result)
+
+    def _call_method(self, prop, args=[], ignore_result=False):
+        return self._link.call_method(self._id, prop, args, ignore_result)
+
+    def _call_method_ignore_return(self, prop, args=[]):
+        return self._link.call_method_ignore_return(self._id, prop, args)
+
+    def _new(self, *args):
+        return self._link.call_new(self._id, args=list(args))
+
+    def _to_js(self):
+        return {
+            "type": "proxy",
+            "id": self._id,
+            "parent_id": self._parent_id,
+        }
+
+    def _get_keys(self):
+        return self._link.get_keys(self._id)
+
+    def __iter__(self):
+        return _ProxyIterator(self)

webgpu/link/websocket.py

@@ -0,0 +1,115 @@
+import asyncio
+import json
+import threading
+
+import websockets
+import websockets.asyncio.client
+
+from .base import LinkBaseAsync
+
+
+class WebsocketLinkBase(LinkBaseAsync):
+    _websocket_thread: threading.Thread
+    _connection: websockets.asyncio.client.ClientConnection
+    _event_is_connected: threading.Event
+    _start_handling_messages: threading.Event
+
+    def __init__(self):
+        super().__init__()
+        self._event_is_connected = threading.Event()
+        self._start_handling_messages = threading.Event()
+        self._send_loop = asyncio.new_event_loop()
+
+        self._websocket_thread = threading.Thread(target=self._connect)
+        self._websocket_thread.start()
+
+    def wait_for_connection(self):
+        self._event_is_connected.wait()
+
+    async def _send_async(self, message):
+        if self._connection:
+            await self._connection.send(message)
+        else:
+            raise Exception("Websocket not connected")
+
+    def _connect(self):
+        raise NotImplementedError
+
+
+class WebsocketLinkClient(WebsocketLinkBase):
+    _url: str
+
+    def __init__(self, url):
+        super().__init__()
+        self._url = url
+
+    def _connect(self):
+        async def start_websocket():
+            async for websocket in websockets.connect(self._url):
+                try:
+                    # print("client connected")
+                    self._connection = websocket
+                    self._event_is_connected.set()
+                    self._start_handling_messages.wait()
+                    async for message in websocket:
+                        self._on_message(message)
+                except websockets.exceptions.ConnectionClosed:
+                    continue
+                except Exception:
+                    break
+
+        try:
+            asyncio.set_event_loop(self._send_loop)
+            self._send_loop.run_until_complete(start_websocket())
+        except Exception:
+            print("closing connection")
+
+
+class WebsocketLinkServer(WebsocketLinkBase):
+    _stop: asyncio.Future
+    _port: int = 0
+
+    def __init__(self):
+        super().__init__()
+        self._stop = self._send_loop.create_future()
+        self._port = 8700
+
+    @property
+    def port(self):
+        return self._port
+
+    async def _websocket_handler(self, websocket, path=""):
+        try:
+            # print("client connected")
+            self._connection = websocket
+            self._event_is_connected.set()
+            async for message in websocket:
+                thread = threading.Thread(target=self._on_message, args=(message,))
+                thread.start()
+        finally:
+            self._connection = None
+
+    def _connect(self):
+        async def start_websocket():
+            while True:
+                try:
+                    async with websockets.serve(
+                        self._websocket_handler, "", self._port, max_queue=128
+                    ):
+                        # print("server running on port", self._port)
+                        await self._stop
+                        break
+                except OSError as e:
+                    self._port += 1
+                except Exception as e:
+                    print("error in websocket server", e)
+
+        try:
+            asyncio.set_event_loop(self._send_loop)
+            self._send_loop.run_until_complete(start_websocket())
+        except Exception as e:
+            print("exception in _start_websocket_server", e)
+        print("stopped websocket")
+
+    def stop(self):
+        self._send_loop.call_soon_threadsafe(self._stop.set_result, None)