{"id":22120203,"url":"https://github.com/joshbrew/onnxruntime-web-fast-style-transfer","last_synced_at":"2025-03-24T06:27:11.876Z","repository":{"id":229906028,"uuid":"777971454","full_name":"joshbrew/onnxruntime-web-fast-style-transfer","owner":"joshbrew","description":"Quick and dirty WebGPU ONNX Runtime test for https://github.com/pytorch/examples/tree/main/fast_neural_style","archived":false,"fork":false,"pushed_at":"2024-03-29T18:05:03.000Z","size":59511,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-01-29T12:14:55.570Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/joshbrew.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-26T20:55:38.000Z","updated_at":"2024-03-26T20:57:12.000Z","dependencies_parsed_at":"2024-12-01T14:31:33.280Z","dependency_job_id":null,"html_url":"https://github.com/joshbrew/onnxruntime-web-fast-style-transfer","commit_stats":null,"previous_names":["joshbrew/onnxruntime-web-fast-style-transfer"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/joshbrew%2Fonnxruntime-web-fast-style-transfer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/joshbrew%2Fonnxruntime-web-fast-style-transfer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/joshbrew%2Fonnxruntime-web-fast-style-transfer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/joshbrew%2Fonnxruntime-web-fast-style-transfer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/joshbrew","download_url":"https://codeload.github.com/joshbrew/onnxruntime-web-fast-style-transfer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":245220136,"owners_count":20579708,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-12-01T14:21:26.853Z","updated_at":"2025-03-24T06:27:11.855Z","avatar_url":"https://github.com/joshbrew.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Jupyter wrapper of [fast-neural-style](https://github.com/pytorch/examples/tree/main/fast_neural_style) example with ONNX converting and a test web app\n\n![result](./testapp/app.png)\n\n## Try it out\n[https://onnx-style-transfer.netlify.app/](https://onnx-style-transfer.netlify.app/)\n\nIt works on 720p-ish PNGs, not sure why it's not working on larger resolutions yet as the buffers are dynamic, and the file format should not matter but that's where it is atm.\n\nOriginal with explanation on usage:\nhttps://github.com/pytorch/examples/tree/main/fast_neural_style\n\nAdditional tweaks based on work for real time Unity and Mobile development done here:\n\nhttps://blog.unity.com/engine-platform/real-time-style-transfer-in-unity-using-deep-neural-networks\n\nand here:\n\nhttps://medium.com/@jamesonthecrow/creating-a-17kb-style-transfer-model-with-layer-pruning-and-quantization-864d7cc53693\n\n\n## Install dependencies\n\npip install -r requirements.txt\n\n## Get or create a dataset.\n\nCreate a folder and subfolder `train/classes` to place all of your training images. The program won't recognize it unless it's under a `classes` subfolder. \n\nThe first images we used were from the recommended COCO 2014 dataset from the original example: https://cocodataset.org/#download\n\nThe `--style-image` defined in [`train.ipynb`](./train.ipynb) is your reference photo. You can use drawings for example, like we did, to adopt its look and feel.\n\nThe neural network will freeze this style in its memory to be able to rapidly reconstruct an image based on an input. It is supposed to be able to perform super resolution on a compressed result as well to speed up the time which we are figuring out how to configure.\n\n## Train\n\nIn `train.ipynb`, adjust these parameters as needed:\n```py\n\ndataset_path = \"train\"\nstyle_image_path = \"test/capture.png\"\nsave_model_dir = \"test\"\nmodel_name = \"model.onnx\"\noutput_image = \"test/output.png\"\n\ncommand = [\n    ...\n    \"--style-weight\", \"1e11\",\n    \"--content-weight\", \"1e5\",\n    \"--epochs\", \"2\",\n    \"--cuda\", \"1\"\n]\n```\nSetting cuda to 0 will force CPU training, which is much slower on a large dataset. Our RTX3070 took 2 hours 20 minutes to train on COCO2014, an 18gb dataset. Speed will vary based on that size.\n\nIf using Mac, see neural_style.py for appropriate commands, uncomment the lines for `args.mps` in neural_style.py (this caused errors on Windows for us in Python 3.11 for whatever reason) and supply the appropriate command in the array according to the [original repository instructions]((https://github.com/pytorch/examples/tree/main/fast_neural_style))\n\nRun `train.ipynb` and let it finish. \n\nFor example we used a still from the movie Mind Game directed by Masaaki Yuasa, which you should go watch right now.\n\n![style](./test/mindgame.jpg)\n\n## Convert\n\nRun `convert.ipynb` to convert the outputted `model.pth` file for Pytorch to a `model.onnx` file for ONNX runtimes using the supplied parameters. We're working on fixing this for `onnxruntime-web` gpu operations. It works in python.\n\nYou can also test the .pth file version first before converting, if you use model_name = \"model.pth\" in the Test.ipynb as shown below.\n\n## Test\n\nTweak parameters as needed:\n\n```py\n\n# Define your paths and options here\nstyle_image_path = \"test/capture.png\"\nsave_model_dir = \"test\"\nmodel_name = \"model.onnx\"\noutput_image = \"test/output2.png\"\n\ncommand2 = [\n    \"python\", \"neural_style.py\", \"eval\",\n    \"--content-image\", style_image_path,\n    \"--model\", save_model_dir+\"/\"+model_name,\n    \"--output-image\", output_image,\n    \"--cuda\", \"1\"\n]\n\n```\n\nRun `test.ipynb`\n\nI am just testing this to see how real-time capable it is for styling a game as a hack, so for example I can feed the frames of a recording in from a game to give the classifier enough context to perform the shader pass.\n\nIn:\n![in](./test/capture.png)\n\nOut:\n![out](./test/baseoutput.png)\n\nNice result but this was on the 128 channel model, which only ran around 3fps on our RTX3070 at 720p on the WebGPU ONNX backend via browser. That's a 7Mb model. We tweaked it and optimized the layer computations and got it down to 41kb without quantizing properly, so we should be able to get much more out of it.\n\nNote the style weight on the smaller model is more finicky so we're tuning it still. You want it pretty high on the large model, we used 1e11, while somewhere between 1e10 and 5e10 is the right setting for the small model. We'll get it!\n\nToo low, does a good reconstruction still:\n![a](./test/mobilenet_lowstyle.webp)\n\nToo high, can't see the scene:\n![a](./test/mobilenet_highstyle.png)\n\nThe border distortions are because of our padding not being totally correct.\n\nWe were able to get 16-20fps at 720p on our mobile version of the model. We're still looking for ways to advance it.\n\n## Run the webapp\n\nCopy the `model.onnx` result you created to testapp/models. Follow the [README](./testapp/README.md) in that subfolder to run the sample. \n\nWe'll be working on optimizing it and seeing if we can get a 30-60fps result. \n\n\n\n## Our Model of Choice (wip)\n\nWe switched all convolutions to Depthwise Separable convolution which should be about 100x less computations, and lowered the number of filters to 9 based on recommendations found online. The resulting model without proper quantization is 41kb, and runs at 16-20fps at 720p on an RTX3070 on the WebGPU ONNX distribution. You can find the .pth or .onnx files in the test/ folder. This should definitely run faster than it does on our hardware.\n\nNote padding shouldn't actually be zero below as it causes artifacts, it's a bug we're fixing.\n```py   \n# Architecture based on: https://medium.com/@jamesonthecrow/creating-a-17kb-style-transfer-model-with-layer-pruning-and-quantization-864d7cc53693\nclass EfficientMobileTransformerNet(torch.nn.Module):\n    def __init__(self):\n        super(EfficientMobileTransformerNet, self).__init__()\n        # Initial convolution layers with less complexity and efficient design\n        self.conv1 = DepthwiseConvLayer(3, 9, kernel_size=9, stride=1, padding=0)\n        self.in1 = torch.nn.InstanceNorm2d(9, affine=True)\n        self.conv2 = DepthwiseConvLayer(9, 9, kernel_size=3, stride=2, padding=0)\n        self.in2 = torch.nn.InstanceNorm2d(9, affine=True)\n        self.conv3 = DepthwiseConvLayer(9, 9, kernel_size=3, stride=2, padding=0)\n        self.in3 = torch.nn.InstanceNorm2d(9, affine=True)\n\n        # A single, more efficient residual block\n        self.resblocks = torch.nn.Sequential(*[OptimizedResidualBlock(9) for _ in range(3)])\n\n        # Upsampling Layers with efficient techniques\n        self.deconv1 = DepthwiseConvLayer(9, 9, kernel_size=3, stride=1, upsample=2, padding=0) #upsample=4\n        self.in4 = torch.nn.InstanceNorm2d(9, affine=True)\n        self.deconv2 = DepthwiseConvLayer(9, 9, kernel_size=9, stride=1, upsample=2, padding=0) #OptimizedUpsampleConvLayer\n        self.in5 = torch.nn.InstanceNorm2d(9, affine=True)\n        self.deconv3 = DepthwiseConvLayer(9, 3, kernel_size=3, stride=1, padding=0)\n       \n        # Non-linearities\n        # ReLU (Rectified Linear Unit) introduces non-linearity, helping the network learn complex patterns.\n        self.relu = torch.nn.ReLU()\n\n    def forward(self, X):\n        # Define the forward pass through the network.\n        \n        # Applies convolutional layers, residual blocks, and upsampling layers in sequence.\n        y = self.relu(self.in1(self.conv1(X)))\n        y = self.relu(self.in2(self.conv2(y)))\n        y = self.relu(self.in3(self.conv3(y)))\n        \n        # Residual blocks do not change the dimensions of their input.\n        y = self.resblocks(y)\n        \n        # Upsampling layers increase the spatial dimensions of the input.\n        y = self.relu(self.in4(self.deconv1(y)))\n        y = self.relu(self.in5(self.deconv2(y)))\n        y = self.deconv3(y)  # Final convolution to produce the output.\n        return y\n\n```\n\n\n## HELP\n\nI'm sharing all results free and making it instructional via these ipynb files as ONNX etc is a bit esoteric but also the cutting edge right now for distributing these kinds of programs. We can make universal shader systems this way, halp me fix bugs and whittle this down to real time capable at least on desktop, probably on mobile too. ","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjoshbrew%2Fonnxruntime-web-fast-style-transfer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fjoshbrew%2Fonnxruntime-web-fast-style-transfer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjoshbrew%2Fonnxruntime-web-fast-style-transfer/lists"}