linter and tests

Author: jfacevedo-google

src/maxdiffusion/generate.py

@@ -196,4 +196,4 @@ def main(argv: Sequence[str]) -> None:
  run(pyconfig.config)
 
 if __name__ == "__main__":
- app.run(main)
+ app.run(main)

src/maxdiffusion/models/train.py

@@ -147,7 +147,8 @@ def train(config):
  weight_dtype = max_utils.get_dtype(config)
  pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
  config.pretrained_model_name_or_path,revision=config.revision, dtype=weight_dtype,
- safety_checker=None, feature_extractor=None, from_pt=config.from_pt
+ safety_checker=None, feature_extractor=None, from_pt=config.from_pt,
+ split_head_dim=config.split_head_dim
  )
 
  noise_scheduler, noise_scheduler_state = FlaxDDPMScheduler.from_pretrained(config.pretrained_model_name_or_path,
@@ -385,4 +386,4 @@ def main(argv: Sequence[str]) -> None:
  validate_train_config(config)
  train(config)
 if __name__ == "__main__":
- app.run(main)
+ app.run(main)

src/maxdiffusion/tests/images/test_2_base.png

@@ -65,30 +65,30 @@ def test_sd21_config(self):
 
  cleanup(output_dir)
 
- # def test_sd_2_base_config(self):
- #  output_dir="train-smoke-test"
- #  train_main([None,os.path.join(THIS_DIR,'..','configs','base_2_base.yml'),
- #  "run_name=sd2_base_smoke_test","max_train_steps=21","dataset_name=lambdalabs/pokemon-blip-captions",
- #  "base_output_directory=gs://maxdiffusion-tests", f"output_dir={output_dir}"])
-
- #  img_url = os.path.join(THIS_DIR,'images','test_2_base.png')
- #  base_image = np.array(Image.open(img_url)).astype(np.uint8)
-
- # pyconfig.initialize([None,os.path.join(THIS_DIR,'..','configs','base2_base_inference.yml'),
- #  f"pretrained_model_name_or_path={output_dir}",
- #  "prompt=A magical castle in the middle of a forest, artistic drawing",
- #  "negative_prompt=purple, red","guidance_scale=7.5",
- #  "num_inference_steps=30","seed=47"])
-
- #  images = generate_run(pyconfig.config)
- # test_image = np.array(images[1]).astype(np.uint8)
- #  ssim_compare = ssim(base_image, test_image,
- #  multichannel=True, channel_axis=-1, data_range=255
- #  )
- #  assert base_image.shape == test_image.shape
- #  assert ssim_compare >=0.70
-
- #  cleanup(output_dir)
+ def test_sd_2_base_config(self):
+ output_dir="train-smoke-test"
+ train_main([None,os.path.join(THIS_DIR,'..','configs','base_2_base.yml'),
+ "run_name=sd2_base_smoke_test","max_train_steps=21","dataset_name=lambdalabs/pokemon-blip-captions",
+ "base_output_directory=gs://maxdiffusion-tests", f"output_dir={output_dir}"])
+
+ img_url = os.path.join(THIS_DIR,'images','test_2_base.png')
+ base_image = np.array(Image.open(img_url)).astype(np.uint8)
+
+  pyconfig.initialize([None,os.path.join(THIS_DIR,'..','configs','base_2_base_inference.yml'),
+ f"pretrained_model_name_or_path={output_dir}",
+ "prompt=A magical castle in the middle of a forest, artistic drawing",
+ "negative_prompt=purple, red","guidance_scale=7.5",
+ "num_inference_steps=30","seed=47"])
+
+ images = generate_run(pyconfig.config)
+  test_image = np.array(images[0]).astype(np.uint8)
+ ssim_compare = ssim(base_image, test_image,
+ multichannel=True, channel_axis=-1, data_range=255
+ )
+ assert base_image.shape == test_image.shape
+ assert ssim_compare >=0.70
+
+ cleanup(output_dir)
 
 if __name__ == '__main__':
  absltest.main()