google-ai-edge / ai-edge-torch

Supporting PyTorch models with the Google AI Edge TFLite runtime.
Apache License 2.0
199 stars 23 forks source link

Refactor decoder of stable diffusion #39

Closed yichunk closed 1 month ago

yichunk commented 1 month ago

Refactor decoder of stable diffusion to use layers of generative API. There are three commits inside the PR.

  1. Add group norm
  2. Add blocks modules and configs used in autoencoder
  3. Refactor decoder of stable diffusion, add model loader and update convert_to_tflite.py pipeline

BUG=b/311216181

haozha111 commented 1 month ago

Great work on making this refactoring!