Test Stable Diffusion text-to-image pipeline for: - PyTorch to CoreML conversion via coremltools - Speed of CoreML runtime across several compute units - Integration with `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py` - Efficacy of the safety_checker - Affi
| 38 | |
| 39 | |
| 40 | class TestStableDiffusionForTextToImage(unittest.TestCase): |
| 41 | """ Test Stable Diffusion text-to-image pipeline for: |
| 42 | |
| 43 | - PyTorch to CoreML conversion via coremltools |
| 44 | - Speed of CoreML runtime across several compute units |
| 45 | - Integration with `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py` |
| 46 | - Efficacy of the safety_checker |
| 47 | - Affinity of the generated image with the original prompt via CLIP score |
| 48 | - The bridge between Python and Swift CLI |
| 49 | - The signal parity of Swift CLI generated image with that of Python CLI |
| 50 | """ |
| 51 | cli_args = None |
| 52 | |
| 53 | @classmethod |
| 54 | def setUpClass(cls): |
| 55 | cls.pytorch_pipe = StableDiffusionPipeline.from_pretrained( |
| 56 | cls.cli_args.model_version, |
| 57 | use_auth_token=True, |
| 58 | ) |
| 59 | |
| 60 | # To be initialized after test_torch_to_coreml_conversion is run |
| 61 | cls.coreml_pipe = None |
| 62 | cls.active_compute_unit = None |
| 63 | |
| 64 | @classmethod |
| 65 | def tearDownClass(cls): |
| 66 | cls.pytorch_pipe = None |
| 67 | cls.coreml_pipe = None |
| 68 | cls.active_compute_unit = None |
| 69 | |
| 70 | def test_torch_to_coreml_conversion(self): |
| 71 | """ Tests: |
| 72 | - PyTorch to CoreML conversion via coremltools |
| 73 | """ |
| 74 | with self.subTest(model="vae_decoder"): |
| 75 | logger.info("Converting vae_decoder") |
| 76 | torch2coreml.convert_vae_decoder(self.pytorch_pipe, self.cli_args) |
| 77 | logger.info("Successfully converted vae_decoder") |
| 78 | |
| 79 | with self.subTest(model="unet"): |
| 80 | logger.info("Converting unet") |
| 81 | torch2coreml.convert_unet(self.pytorch_pipe, self.cli_args) |
| 82 | logger.info("Successfully converted unet") |
| 83 | |
| 84 | with self.subTest(model="text_encoder"): |
| 85 | logger.info("Converting text_encoder") |
| 86 | torch2coreml.convert_text_encoder(self.pytorch_pipe, self.cli_args) |
| 87 | logger.info("Successfully converted text_encoder") |
| 88 | |
| 89 | with self.subTest(model="safety_checker"): |
| 90 | logger.info("Converting safety_checker") |
| 91 | torch2coreml.convert_safety_checker(self.pytorch_pipe, |
| 92 | self.cli_args) |
| 93 | logger.info("Successfully converted safety_checker") |
| 94 | |
| 95 | def test_end_to_end_image_generation_speed(self): |
| 96 | """ Tests: |
| 97 | - Speed of CoreML runtime across several compute units |
no outgoing calls
no test coverage detected