import ast from copy import deepcopy from contextlib import nullcontext import io import json from pathlib import Path import re import sys import tempfile from types import SimpleNamespace import unittest from unittest.mock import Mock import gradio as gr import numpy as np from PIL import Image sys.path.insert(0,str(Path(__file__).resolve().parents[1])) from generation_metadata import JSON_KEY, generation_record, image_with_metadata, infotext def record(family="anima",**kwargs): values=dict(family=family,model="anima-turbo-v1.1_INT8.safetensors" if family=="anima" else "waiIllustriousSDXL_v170.safetensors", input_prompt=", 日本語\nblue eyes",input_negative_prompt="bad quality", prompt="red hair, 日本語\nblue eyes",negative_prompt="bad quality", width=16,height=16,steps=8,cfg=1.0,seed=9223372036854775806, sampler="er_sde",scheduler="simple",vae="qwen_image_vae.safetensors",vae_source="external", loras=[],embeddings=[],comfyui_commit="pinned",int8=family=="anima") values.update(kwargs) return generation_record(**values) class TestGenerationMetadata(unittest.TestCase): def test_png_both_formats_and_pixels_roundtrip(self): pixels=np.arange(16*16*3,dtype=np.uint8).reshape(16,16,3) metadata=record() image=image_with_metadata(pixels,metadata) from gradio.processing_utils import encode_pil_to_bytes saved=Image.open(io.BytesIO(encode_pil_to_bytes(image,format="png"))) np.testing.assert_array_equal(np.array(saved),pixels) self.assertEqual(json.loads(saved.info[JSON_KEY]),metadata) self.assertIn("日本語\nblue eyes\nNegative prompt: bad quality",saved.info["parameters"]) self.assertIn("Seed: 9223372036854775806",saved.info["parameters"]) self.assertIn("RNG: CPU",saved.info["parameters"]) self.assertEqual(json.loads(saved.info[JSON_KEY])["seed_string"],"9223372036854775806") def test_ill_embedded_external_vae_and_lora_fields(self): for source,vae,expected in (("checkpoint","Checkpoint内蔵VAE","VAE: None"), ("external","sdxl_vae.safetensors","VAE: sdxl_vae.safetensors")): metadata=record("illustrious",vae_source=source,vae=vae,sampler="euler_ancestral",scheduler="normal", loras=[dict(slot=1,name='名前,引用".safetensors',strength_model=.7,strength_clip=0.0)], embeddings=["negativeXL_D.safetensors"],negative_prompt="embedding:negativeXL_D.safetensors, bad quality") text=infotext(metadata) self.assertIn("Sampler: Euler a",text); self.assertIn("Schedule type: Normal",text) self.assertIn("Clip skip: 2",text); self.assertIn(expected,text) self.assertIn('LoRA 1: "名前,引用\\\".safetensors"',text) self.assertIn("LoRA 1 strength: 0.7",text) self.assertEqual(metadata["clip_stop_at_layer"],-2) def test_sampler_native_ids_not_lost(self): metadata=record(sampler="er_sde",scheduler="simple") self.assertIn("Sampler: er_sde",infotext(metadata)) self.assertIn("ComfyUI sampler: er_sde",infotext(metadata)) self.assertNotIn("Clip skip:",infotext(metadata)) self.assertEqual(metadata["sampler"],"er_sde") def test_lora_tags_in_png_positive_only_preserve_native_record_and_pixels(self): from gradio.processing_utils import encode_pil_to_bytes for family in ("anima", "illustrious"): metadata=record(family,loras=[ dict(slot=1,name="subdir/日本語.v2.safetensors",strength_model=.7,strength_clip=0.0), dict(slot=4,name="Turbo-ANIMA-v2.safetensors",strength_model=0.0,strength_clip=0.0), ]) before=deepcopy(metadata) pixels=np.full((16,16,3),128,dtype=np.uint8) png=Image.open(io.BytesIO(encode_pil_to_bytes(image_with_metadata(pixels,metadata),format="png"))) positive,remaining=png.info["parameters"].split("\nNegative prompt: ",1) self.assertEqual(positive,metadata["prompt"]+" ") self.assertEqual(remaining.splitlines()[0],metadata["negative_prompt"]) self.assertEqual(json.loads(png.info[JSON_KEY]),before) self.assertEqual(metadata,before) np.testing.assert_array_equal(np.asarray(png),pixels) def test_cache_separates_same_pixels_different_metadata(self): pixels=np.zeros((16,16,3),dtype=np.uint8) with tempfile.TemporaryDirectory() as directory: for streaming in (False,True): with gr.Blocks(analytics_enabled=False): component=gr.Image(format="png",streaming=streaming) component.GRADIO_CACHE=directory first=component.postprocess(image_with_metadata(pixels,record(seed=123))) second=component.postprocess(image_with_metadata(pixels,record(seed=456))) self.assertNotEqual(first.path,second.path) with Image.open(first.path) as image: self.assertEqual(json.loads(image.info[JSON_KEY])["seed"],123) with Image.open(second.path) as image: self.assertEqual(json.loads(image.info[JSON_KEY])["seed"],456) def test_metadata_is_per_job_and_drops_unrelated_info(self): loras=[dict(slot=1,name="a",strength_model=.5,strength_clip=0)] metadata=record(loras=loras) loras[0]["name"]="changed" self.assertEqual(metadata["loras"][0]["name"],"a") original=Image.new("RGB",(16,16)); original.info["unrelated-secret"]="do not propagate" image=image_with_metadata(original,metadata) self.assertNotIn("unrelated-secret",image.info) self.assertIn("unrelated-secret",original.info) def test_core_records_executed_values_not_requested_defaults(self): source=ast.parse((Path(__file__).resolve().parents[1]/"app.py").read_text(encoding="utf-8")) function=deepcopy(next(n for n in source.body if isinstance(n,ast.FunctionDef) and n.name=="_generate_image_core")) function.decorator_list=[] pixels=np.zeros((512,1024,3),dtype=np.uint8) tensor=Mock() tensor.mul.return_value=tensor; tensor.clamp_.return_value=tensor tensor.byte.return_value=tensor; tensor.cpu.return_value=tensor; tensor.numpy.return_value=pixels for family in ("anima","illustrious"): name="actual_INT8.safetensors" if family=="anima" else "wai.safetensors" lora_loader=SimpleNamespace(load_lora_model_only=Mock(return_value=("patched",))) env=dict(np=np,gr=SimpleNamespace(Progress=lambda **kw:None,Error=ValueError), torch=SimpleNamespace(nan_to_num=lambda image,**kw:image), random=SimpleNamespace(randint=lambda *a:9223372036854775806), get_value_at_index=lambda obj,i:obj[i], model_components=lambda *args:("model","clip","vae",family,"qwen_image_vae.safetensors" if family=="anima" else "Checkpoint内蔵VAE"), lora_selected=lambda enabled,name:enabled and name!="---", FAMILY_LORA_NAMES={family:{"details","turbo"}},lora_loader=lora_loader, normalize_prompt_text=lambda text,rules:text.replace("","red hair"), POSITIVE_PROMPT_REPLACE_RULES=[],NEGATIVE_PROMPT_REPLACE_RULES=[], apply_negative_embeddings=lambda text,*a:("embedding:neg.safetensors, "+text,["neg.safetensors"]), text_encode=SimpleNamespace(encode=lambda **kw:(kw["text"],)), empty_latent=SimpleNamespace(generate=lambda **kw:({"samples":"latent"},)), comfy_sample=SimpleNamespace(fix_empty_latent_channels=lambda model,latent,*a:latent, prepare_noise=lambda *a:None,sample=lambda **kw:"sampled"), vae_decode=SimpleNamespace(decode=lambda **kw:([tensor],)), reference_model=lambda model,*a:nullcontext((model,None)),COMFYUI_PATH="fixture", STATIC_INFO="info",gpu_stats=lambda:"gpu",generation_record=generation_record, CHECKPOINT_VAE_OPTION="embedded",COMFYUI_COMMIT="pinned",INT8_ANIMA_MODELS={"actual_INT8.safetensors"}) exec(compile(ast.Module(body=[function],type_ignores=[]),"app.py","exec"),env) output=env["_generate_image_core"](name,"","bad",1025,509,-1,999,99, "er_sde","simple",True,"details",3.5,False,"off",1,True,"---",1,True,"turbo",-1, ["neg.safetensors"],None,lambda *a,**kw:None) meta=output[4] self.assertEqual((meta["width"],meta["height"],meta["steps"],meta["cfg"]),(1024,512,60,8.0)) self.assertEqual(meta["seed"],9223372036854775806) self.assertEqual(meta["model"],name) self.assertEqual(meta["input_prompt"],"") self.assertEqual(meta["prompt"],"red hair") self.assertEqual(meta["negative_prompt"],"embedding:neg.safetensors, bad") self.assertEqual([(r["slot"],r["strength_model"]) for r in meta["loras"]],[(1,2.0),(4,0.0)]) self.assertEqual(meta["vae_source"],"external" if family=="anima" else "checkpoint") self.assertEqual(meta["int8"],family=="anima") if __name__=="__main__": unittest.main()