{"id":15692509,"url":"https://github.com/sayakpaul/xla-benchmark-sd","last_synced_at":"2025-05-08T02:41:54.853Z","repository":{"id":106650383,"uuid":"607069550","full_name":"sayakpaul/xla-benchmark-sd","owner":"sayakpaul","description":"Provides code to serialize the different models involved in Stable Diffusion as SavedModels and to compile them with XLA.","archived":false,"fork":false,"pushed_at":"2023-02-27T10:57:50.000Z","size":28,"stargazers_count":7,"open_issues_count":0,"forks_count":3,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-31T16:52:14.004Z","etag":null,"topics":["compilers-for-ml","image-generation","keras","kerascv","stable-diffusion","tensorflow","xla"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sayakpaul.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"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":"2023-02-27T08:46:59.000Z","updated_at":"2023-12-30T13:36:04.000Z","dependencies_parsed_at":null,"dependency_job_id":"b1498e86-5afb-4ef0-aa58-64261d4cde5d","html_url":"https://github.com/sayakpaul/xla-benchmark-sd","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fxla-benchmark-sd","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fxla-benchmark-sd/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fxla-benchmark-sd/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sayakpaul%2Fxla-benchmark-sd/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sayakpaul","download_url":"https://codeload.github.com/sayakpaul/xla-benchmark-sd/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":252989670,"owners_count":21836662,"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":["compilers-for-ml","image-generation","keras","kerascv","stable-diffusion","tensorflow","xla"],"created_at":"2024-10-03T18:34:30.520Z","updated_at":"2025-05-08T02:41:54.837Z","avatar_url":"https://github.com/sayakpaul.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# XLA compilation of Stable Diffusion in TensorFlow\n\n**_TL;DR: You can XLA-compile the isolated parts of Stable Diffusion in TensorFlow and achieve 17% speedup boost in inference._**\n\nThis repository provides code to serialize the different models involved in Stable Diffusion as SavedModels and to compile them with XLA. As result of the XLA-compiled concrete functions, we can obtain a good amount of speedup in the inference process.\n\nWe use the Stable Diffusion model [shipped](https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/) from KerasCV.\n\n**Table of content**\n\n* [Results](#results)\n* [Steps](#steps)\n* [Running the benchmark](#running-the-benchmark)\n* [Benchmark details](#details-of-the-benchmark)\n* [Gotchas](#gotchas)\n* [Acknowledgements](#acknowledgements)\n\n## Results \n\n* KerasCV with XLA: 12.40 seconds\n* SavedModels with XLA 10.29 seconds\n* SavedModels without XLA 13.69 seconds\n\n**_~25% w.r.t non-XLA SavedModel \u0026 ~17% w.r.t KerasCV._** \n\n## Steps \n\nWe first isolate the sub-models involved in Stable Diffusion and serialize them as\nstand-alone SavedModels:\n\n* Text encoder\n* Diffusion model aka UNet\n* Decoder\n\nThe SavedModel also includes their respective computations. For example SavedModel of the text encoder includes the processing of prompt context and the unconditional context. Similarly, SavedModel of the UNet includes the computations for the diffusion \nprocess. \n\nFor the serialization, just run `serialize_savedmodels.py`. \n\nOnce the SavedModels are generated, we load them as concrete functions and XLA-compile them before running inference. We include the complete code for this in `benchmark.py`. \n\n## Running the benchmark\n\nFor running the KerasCV benchmark:\n\n```bash\npython benchmark.py --kerascv --jit_compile\n```\n\nFor running with SavedModels (**without** XLA):\n\n```bash\npython benchmark.py \n```\n\nFor running with SavedModels (**with** XLA):\n\n```bash\npython benchmark.py --jit_compile\n```\n\n## Details of the benchmark\n\nThe benchmarks were run on an `a2-highgpu-1g` [instance](https://cloud.google.com/compute/docs/gpus#a100-gpus). \n\n## Gotchas\n\n* The text encoder cannot be XLA-compiled. See [this issue](https://github.com/tensorflow/tensorflow/issues/59818) for more details.\n* For making the SavedModels XLA-compitable, we fix the number of images that can be generated per prompt. Otherwise, it doesn't become a compile-time constant which \nmakes it XLA-incompatible.\n\n## Acknowledgements\n\nThanks to the ML Developer Programs' team at Google for providing GCP credit support.","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fxla-benchmark-sd","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsayakpaul%2Fxla-benchmark-sd","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsayakpaul%2Fxla-benchmark-sd/lists"}