File size: 4,383 Bytes
3173362 f4ade94 33644f7 df262ae f4ade94 3173362 f4ade94 3173362 df262ae f4ade94 3173362 33644f7 3173362 33644f7 3173362 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | 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()
|