updated.

Author: tanyuqian

examples/text_to_image/main.py

@@ -122,7 +122,7 @@ def main(dataset_name='lambdalabs/naruto-blip-captions',
          model_name_or_path='stabilityai/stable-diffusion-2-1-base',
          init_ckpt_dir='./stable-diffusion-2-1-base',
          n_model_shards=1,
-         n_epochs=3,
+         n_epochs=8,
          global_batch_size=8,
          per_device_batch_size=1,
          learning_rate=1e-5,
@@ -156,13 +156,13 @@ def main(dataset_name='lambdalabs/naruto-blip-captions',
     text_encoder = FlaxCLIPTextModel(
         config=CLIPTextConfig.from_pretrained(
             model_name_or_path, subfolder="text_encoder"),
-        dtype=jnp.float16, _do_init=False)
+        dtype=jnp.float32, _do_init=False)
     vae = FlaxAutoencoderKL.from_config(
         config=FlaxAutoencoderKL.load_config(
-            model_name_or_path, subfolder='vae'), dtype=jnp.float16)
+            model_name_or_path, subfolder='vae'), dtype=jnp.float32)
     unet = FlaxUNet2DConditionModel.from_config(
         config=FlaxUNet2DConditionModel.load_config(
-            model_name_or_path, subfolder='unet'), dtype=jnp.float16)
+            model_name_or_path, subfolder='unet'), dtype=jnp.float32)
     noise_scheduler, noise_scheduler_state = \
         FlaxPNDMScheduler.from_pretrained(
             model_name_or_path, subfolder='scheduler')

examples/text_to_image/save_init_ckpt.py

@@ -1,6 +1,5 @@
 import fire
 import jax
-import jax.numpy as jnp
 from transformers import FlaxCLIPTextModel
 from diffusers import FlaxAutoencoderKL, FlaxUNet2DConditionModel
 from redco import Deployer
@@ -14,14 +13,11 @@ def main(model_name_or_path='stabilityai/stable-diffusion-2-1-base'):
 
     with jax.default_device(jax.local_devices(backend='cpu')[0]):
         text_encoder = FlaxCLIPTextModel.from_pretrained(
-            model_name_or_path,
-            subfolder="text_encoder", from_pt=True, dtype=jnp.float16)
+            model_name_or_path, subfolder="text_encoder", from_pt=True)
         vae, vae_params = FlaxAutoencoderKL.from_pretrained(
-            model_name_or_path,
-            subfolder="vae", from_pt=True, dtype=jnp.float16)
+            model_name_or_path, subfolder="vae", from_pt=True)
         unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
-            model_name_or_path,
-            subfolder="unet", from_pt=True, dtype=jnp.float32)
+            model_name_or_path, subfolder="unet", from_pt=True)
         params = {
             'text_encoder': text_encoder.params,
             'unet': unet_params,