File size: 3,537 Bytes
e9190ff | 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 | import json
import tempfile
import unittest
from pathlib import Path
from PIL import Image
from scripts.integrity import sha256
from scripts.editing_protocol import image_diagnostics
from scripts.editing_report import validate_run
class EditingReportTests(unittest.TestCase):
def fixture(self, root):
image = Image.new('RGBA', (32, 32), 'white')
image.save(root/'input.png')
image.save(root/'result.png')
case = dict(id='test', prompt='Keep unchanged', source_seed=42,
input_file='input.png', input_sha256=sha256(root/'input.png'))
(root/'cases.json').write_text(json.dumps([case]))
env = dict(protocol='editing-v2', seeds=[1000042], steps=40, resolution=32,
kv_cache=True, vae_tiling=False, offload='model',
before_load_memory={'allocated_bytes': 0},
cases_sha256=sha256(root/'cases.json'))
(root/'environment.json').write_text(json.dumps(env))
row = dict(protocol='editing-v2', case_id='test', prompt=case['prompt'],
seed=1000042, source_seed=42, steps=40, output_resolution=32,
width=32, height=32, cfg=1., kv_cache=True, vae_tiling=False,
offload='model', input_file='input.png', input_sha256=case['input_sha256'],
image='result.png', image_sha256=sha256(root/'result.png'),
**image_diagnostics(image))
self.write_rows(root, [row])
return row
def write_rows(self, root, rows):
(root/'records.jsonl').write_text(''.join(json.dumps(r)+'\n' for r in rows))
(root/'COMPLETE.json').write_text(json.dumps(dict(records=len(rows),
records_sha256=sha256(root/'records.jsonl'))))
def test_valid_run_and_rejects_duplicate_missing_or_wrong_seed(self):
with tempfile.TemporaryDirectory() as d:
root = Path(d); row = self.fixture(root)
self.assertEqual(len(validate_run(root, root)[1]), 1)
for rows in ([row, row], [], [dict(row, seed=42)]):
self.write_rows(root, rows)
with self.assertRaises(ValueError): validate_run(root, root)
def test_rejects_diagnostics_that_claim_white_is_transparent(self):
with tempfile.TemporaryDirectory() as d:
root = Path(d); row = self.fixture(root)
self.write_rows(root, [dict(row, alpha_transparent_fraction=1.)])
with self.assertRaisesRegex(ValueError, 'diagnostic'): validate_run(root, root)
def test_rejects_incomplete_run_and_mutated_input(self):
with tempfile.TemporaryDirectory() as d:
root = Path(d); self.fixture(root)
(root/'COMPLETE.json').unlink()
with self.assertRaises((ValueError, FileNotFoundError)): validate_run(root, root)
self.fixture(root)
Image.new('RGB', (32, 32), 'red').save(root/'input.png')
with self.assertRaisesRegex(ValueError, 'Input'): validate_run(root, root)
def test_rejects_png_dimensions_even_if_diagnostics_agree(self):
with tempfile.TemporaryDirectory() as d:
root = Path(d); row = self.fixture(root)
image = Image.new('RGBA', (64, 32), 'white')
image.save(root/'result.png')
row.update(image_sha256=sha256(root/'result.png'), **image_diagnostics(image))
self.write_rows(root, [row])
with self.assertRaisesRegex(ValueError, 'dimensions'): validate_run(root, root)
|