| import unittest |
| import os |
| import tempfile |
| import json |
| import io |
| import types |
| from PIL import Image |
| from pathlib import Path |
| from unittest.mock import patch |
|
|
| from qwen21.handler import cached_model_path, handler, sampling |
|
|
|
|
| class SamplingTests(unittest.TestCase): |
| def test_native_rgba_photo_is_saved_as_rgb_without_recoloring(self): |
| source = io.BytesIO() |
| Image.new('RGB', (64, 64), 'white').save(source, 'PNG') |
| photo = Image.new('RGBA', (64, 64), (123, 145, 167, 254)) |
| pipe = unittest.mock.MagicMock() |
| pipe.return_value.images = [photo] |
| pipe.scheduler.timesteps = range(40) |
| torch = types.SimpleNamespace( |
| cuda=types.SimpleNamespace(synchronize=lambda: None, |
| reset_peak_memory_stats=lambda: None, get_device_name=lambda: '48 GB test GPU', |
| max_memory_reserved=lambda: 20 * 1024**3), |
| Generator=lambda device: types.SimpleNamespace(manual_seed=lambda seed: None)) |
| with patch.dict('sys.modules', {'torch': torch}), \ |
| patch('qwen21.handler.pipeline', return_value=pipe), \ |
| patch('qwen21.handler.requests.get') as get, \ |
| patch('qwen21.handler.upload', return_value='https://result.png') as save, \ |
| patch.dict(os.environ, {'QWEN21_MODEL_REVISION': 'pinned'}): |
| get.return_value.content = source.getvalue() |
| result = handler({'input': {'task': 'edit', 'jobs': [{ |
| 'image_url': 'https://source.png', 'prompt': 'A complete photograph.'}]}}) |
| self.assertTrue(result['results'][0]['ok']) |
| self.assertEqual(result['results'][0]['native_alpha_range'], (254, 254)) |
| saved = save.call_args.args[0] |
| self.assertEqual(saved.mode, 'RGB') |
| self.assertEqual(saved.getpixel((0, 0)), (123, 145, 167)) |
|
|
| def test_only_the_pinned_runpod_snapshot_is_used(self): |
| with tempfile.TemporaryDirectory() as root, patch.dict(os.environ, { |
| 'HF_HUB_CACHE': root, 'QWEN21_MODEL_REPO': 'Qwen/example', |
| 'QWEN21_MODEL_REVISION': 'pinned'}): |
| with self.assertRaises(FileNotFoundError): |
| cached_model_path() |
| wrong = Path(root) / 'models--Qwen--example/snapshots/wrong' |
| wrong.mkdir(parents=True) |
| (wrong / 'model_index.json').write_text('{}') |
| with self.assertRaises(FileNotFoundError): |
| cached_model_path() |
| expected = wrong.parent / 'pinned' |
| expected.mkdir() |
| (expected / 'model_index.json').write_text('{}') |
| self.assertEqual(cached_model_path(), str(expected)) |
|
|
| def test_official_base_recipe_uses_prefix_cache_without_2511_guidance(self): |
| args = sampling({'prompt': 'Preserve the photographed bag.', 'size': '1024*1536'}) |
| self.assertEqual((args['width'], args['height'], args['num_inference_steps']), (1024, 1536, 40)) |
| self.assertEqual(args['true_cfg_scale'], 1.0) |
| self.assertTrue(args['use_kv_cache']) |
| self.assertEqual(args['output_resolution'], 1024) |
| self.assertNotIn('guidance_scale', args) |
| self.assertNotIn('sigmas', args) |
|
|
| def test_sampling_matches_the_pinned_official_pipeline_signature(self): |
| root = Path(__file__).parents[1] |
| contract = json.loads((root / 'tests/pipeline_contract.json').read_text()) |
| self.assertIn(contract['diffusers_commit'], (root / 'requirements.txt').read_text()) |
| self.assertLessEqual(set(sampling({'prompt': 'Keep the actual bag.'})), |
| set(contract['call_parameters'])) |
|
|
| def test_invalid_controls_fail_before_loading_a_model(self): |
| with patch('qwen21.handler.pipeline') as load: |
| result = handler({'input': {'task': 'edit', 'jobs': [ |
| {'prompt': 'Keep it.', 'image_url': 'https://source.jpg', 'guidance_scale': 4}]}}) |
| load.assert_not_called() |
| self.assertFalse(result['results'][0]['ok']) |
|
|
| def test_resolution_steps_and_empty_prompt_are_validated(self): |
| for values in ({'size': '1025*1536'}, {'num_inference_steps': 200}, |
| {'true_cfg_scale': 0}, {'prompt': ''}): |
| with self.subTest(values=values), self.assertRaises(ValueError): |
| sampling({'prompt': 'Keep it.', **values}) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
