Spaces:
Running
Running
| import unittest | |
| from unittest.mock import patch | |
| import numpy as np | |
| import torch | |
| from acestep.core.generation.handler.diffusion import DiffusionMixin | |
| class _Host(DiffusionMixin): | |
| def __init__(self, device: str = "cpu", dtype: torch.dtype = torch.float32): | |
| self.mlx_decoder = object() | |
| self.device = device | |
| self.dtype = dtype | |
| class _IterableTimesteps: | |
| def __init__(self, values): | |
| self._values = values | |
| def __iter__(self): | |
| return iter(self._values) | |
| class DiffusionMixinTests(unittest.TestCase): | |
| def test_mlx_run_diffusion_converts_inputs_and_outputs_tensor(self): | |
| host = _Host(dtype=torch.float16) | |
| encoder_hidden_states = torch.randn(2, 4, 8, dtype=torch.float64) | |
| encoder_attention_mask = torch.ones(2, 4, dtype=torch.int64) | |
| context_latents = torch.randn(2, 16, 8, dtype=torch.float64) | |
| src_latents = torch.zeros(2, 3, 5, dtype=torch.float32) | |
| timesteps = torch.tensor([1.0, 0.5], dtype=torch.float32) | |
| non_cover_hidden = torch.randn(2, 4, 8, dtype=torch.float64) | |
| non_cover_mask = torch.ones(2, 4, dtype=torch.int64) | |
| non_cover_context = torch.randn(2, 16, 8, dtype=torch.float64) | |
| fake_target = np.ones((2, 3, 5), dtype=np.float32) | |
| def _fake_generate(**kwargs): | |
| self.assertIs(kwargs["mlx_decoder"], host.mlx_decoder) | |
| self.assertEqual(kwargs["src_latents_shape"], (2, 3, 5)) | |
| self.assertEqual(kwargs["timesteps"], [1.0, 0.5]) | |
| self.assertEqual(kwargs["infer_method"], "sde") | |
| self.assertEqual(kwargs["shift"], 2.0) | |
| self.assertEqual(kwargs["audio_cover_strength"], 0.6) | |
| self.assertEqual(kwargs["encoder_hidden_states_np"].dtype, np.float32) | |
| self.assertEqual(kwargs["context_latents_np"].dtype, np.float32) | |
| self.assertEqual(kwargs["encoder_hidden_states_non_cover_np"].dtype, np.float32) | |
| self.assertEqual(kwargs["context_latents_non_cover_np"].dtype, np.float32) | |
| return {"target_latents": fake_target, "time_costs": {"diffusion_time_cost": 1.2}} | |
| with patch("acestep.core.generation.handler.diffusion.mlx_generate_diffusion", side_effect=_fake_generate): | |
| result = host._mlx_run_diffusion( | |
| encoder_hidden_states=encoder_hidden_states, | |
| encoder_attention_mask=encoder_attention_mask, | |
| context_latents=context_latents, | |
| src_latents=src_latents, | |
| seed=123, | |
| infer_method="sde", | |
| shift=2.0, | |
| timesteps=timesteps, | |
| audio_cover_strength=0.6, | |
| encoder_hidden_states_non_cover=non_cover_hidden, | |
| encoder_attention_mask_non_cover=non_cover_mask, | |
| context_latents_non_cover=non_cover_context, | |
| ) | |
| self.assertIn("target_latents", result) | |
| self.assertIn("time_costs", result) | |
| self.assertEqual(result["time_costs"]["diffusion_time_cost"], 1.2) | |
| self.assertEqual(result["target_latents"].dtype, torch.float16) | |
| self.assertEqual(result["target_latents"].device.type, "cpu") | |
| self.assertTrue(torch.allclose(result["target_latents"], torch.ones_like(result["target_latents"]))) | |
| def test_mlx_run_diffusion_handles_optional_and_iterable_timesteps(self): | |
| host = _Host(dtype=torch.float32) | |
| encoder_hidden_states = torch.randn(1, 2, 3, dtype=torch.float32) | |
| encoder_attention_mask = torch.ones(1, 2, dtype=torch.int64) | |
| context_latents = torch.randn(1, 4, 3, dtype=torch.float32) | |
| src_latents = torch.zeros(1, 2, 3, dtype=torch.float32) | |
| timesteps = _IterableTimesteps([0.9, 0.8, 0.7]) | |
| def _fake_generate(**kwargs): | |
| self.assertEqual(kwargs["timesteps"], [0.9, 0.8, 0.7]) | |
| self.assertIsNone(kwargs["encoder_hidden_states_non_cover_np"]) | |
| self.assertIsNone(kwargs["context_latents_non_cover_np"]) | |
| return {"target_latents": np.zeros((1, 2, 3), dtype=np.float32), "time_costs": {}} | |
| with patch("acestep.core.generation.handler.diffusion.mlx_generate_diffusion", side_effect=_fake_generate): | |
| result = host._mlx_run_diffusion( | |
| encoder_hidden_states=encoder_hidden_states, | |
| encoder_attention_mask=encoder_attention_mask, | |
| context_latents=context_latents, | |
| src_latents=src_latents, | |
| seed=1, | |
| timesteps=timesteps, | |
| ) | |
| self.assertEqual(tuple(result["target_latents"].shape), (1, 2, 3)) | |
| self.assertEqual(result["target_latents"].dtype, torch.float32) | |
| def test_mlx_run_diffusion_rejects_invalid_infer_method(self): | |
| host = _Host() | |
| x = torch.randn(1, 2, 3) | |
| with self.assertRaises(ValueError): | |
| host._mlx_run_diffusion( | |
| encoder_hidden_states=x, | |
| encoder_attention_mask=torch.ones(1, 2, dtype=torch.int64), | |
| context_latents=torch.randn(1, 4, 3), | |
| src_latents=torch.randn(1, 2, 3), | |
| seed=1, | |
| infer_method="bad", | |
| ) | |
| def test_mlx_run_diffusion_rejects_non_iterable_timesteps(self): | |
| host = _Host() | |
| x = torch.randn(1, 2, 3) | |
| with self.assertRaises(TypeError): | |
| host._mlx_run_diffusion( | |
| encoder_hidden_states=x, | |
| encoder_attention_mask=torch.ones(1, 2, dtype=torch.int64), | |
| context_latents=torch.randn(1, 4, 3), | |
| src_latents=torch.randn(1, 2, 3), | |
| seed=1, | |
| timesteps=123, | |
| ) | |
| def test_mlx_run_diffusion_rejects_batch_mismatch(self): | |
| host = _Host() | |
| with self.assertRaises(ValueError): | |
| host._mlx_run_diffusion( | |
| encoder_hidden_states=torch.randn(2, 2, 3), | |
| encoder_attention_mask=torch.ones(2, 2, dtype=torch.int64), | |
| context_latents=torch.randn(1, 4, 3), | |
| src_latents=torch.randn(2, 2, 3), | |
| seed=1, | |
| ) | |
| def test_mlx_run_diffusion_requires_host_attributes(self): | |
| class _BrokenHost(DiffusionMixin): | |
| pass | |
| host = _BrokenHost() | |
| x = torch.randn(1, 2, 3) | |
| with self.assertRaises(AttributeError): | |
| host._mlx_run_diffusion( | |
| encoder_hidden_states=x, | |
| encoder_attention_mask=torch.ones(1, 2, dtype=torch.int64), | |
| context_latents=torch.randn(1, 4, 3), | |
| src_latents=torch.randn(1, 2, 3), | |
| seed=1, | |
| ) | |
| if __name__ == "__main__": | |
| unittest.main() | |