Add UNet model config and refactor diffusion model.

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

Closed yichunk closed 3 weeks ago

yichunk commented 3 weeks ago

Add UNet model config and refactor diffusion model. convert_to_tflite.py of stable diffusion works, and the pipeline generates reasonable result.

BUG=b/311216181