search

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

Add 2D blocks used in diffusion model of stable diffusion. #50

Closed yichunk closed 3 weeks ago

yichunk commented 4 weeks ago

Add 2D blocks used in diffusion model of stable diffusion.

  1. FeedForwardBlock2D: used in TransformerBlock2D
  2. CrossAttentionBlock2D: used in TransformerBlock2D
  3. TransformerBlock2D: used in DownEncoder2DBlock and SkipUpDecoderBlock2D
  4. DownEncoder2DBlock: used in encoders of diffusion model
  5. SkipUpDecoderBlock2D: used in decoders of diffusion model

BUG=b/311216181