feat: snapshot product suite reference assets
This commit is contained in:
@@ -75,6 +75,7 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
||||
"recovery_action",
|
||||
"generation_round_key",
|
||||
"generation_slot_index",
|
||||
"reference_asset_ids",
|
||||
}.issubset(jobs_columns)
|
||||
)
|
||||
indexes = {
|
||||
@@ -843,6 +844,48 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
||||
self.assertTrue(
|
||||
all(job.generation_round_key is None for job in legacy_jobs)
|
||||
)
|
||||
self.assertTrue(
|
||||
all(job.reference_asset_ids is None for job in legacy_jobs)
|
||||
)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_create_job_freezes_valid_reference_asset_ids(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
config = self._config(temp_dir)
|
||||
db.init_db(config["db_path"])
|
||||
project = image_studio.create_or_get_project(
|
||||
account_alias="店铺",
|
||||
account_slug="shop",
|
||||
item_id="51100639510",
|
||||
path=config["db_path"],
|
||||
)
|
||||
source = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
path=config["db_path"],
|
||||
)
|
||||
reference = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
path=config["db_path"],
|
||||
)
|
||||
job = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
reference_asset_ids=[reference.id],
|
||||
path=config["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual("[%d]" % reference.id, job.reference_asset_ids)
|
||||
self.assertEqual([reference.id], image_studio.job_reference_asset_ids(job))
|
||||
with self.assertRaisesRegex(db.DbError, "不能包含主图"):
|
||||
image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
reference_asset_ids=[source.id],
|
||||
path=config["db_path"],
|
||||
)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
@@ -246,6 +246,89 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_job_reference_snapshot_submits_ordered_images(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
reference_path = os.path.join(temp_dir, "reference.png")
|
||||
with open(reference_path, "wb") as fh:
|
||||
fh.write(self._png_bytes())
|
||||
reference = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
local_path=reference_path,
|
||||
source_order=2,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
job = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
reference_asset_ids=[reference.id],
|
||||
prompt="多图提示词",
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
submitted = []
|
||||
|
||||
def fake_submit(method, url, api_key, **kwargs):
|
||||
submitted.append(kwargs["payload"])
|
||||
return {"task_id": "multi-image-task", "status": "queued"}
|
||||
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry", side_effect=fake_submit), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._cmhub_call_once",
|
||||
return_value={
|
||||
"task_id": "multi-image-task",
|
||||
"status": "succeeded",
|
||||
"result": {"image_url": "https://cdn.example.com/multi.png"},
|
||||
},
|
||||
), \
|
||||
mock.patch(
|
||||
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
|
||||
return_value=(self._png_bytes(), 0.1),
|
||||
):
|
||||
summary = image_studio_generation.run_jobs(
|
||||
[job],
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(1, summary["success"])
|
||||
self.assertEqual(2, len(submitted[0]["images"]))
|
||||
self.assertNotIn("image_base64", submitted[0])
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_missing_reference_snapshot_fails_without_submitting(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
reference = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
local_path=os.path.join(temp_dir, "missing-reference.png"),
|
||||
source_order=2,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
job = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
reference_asset_ids=[reference.id],
|
||||
prompt="多图提示词",
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
|
||||
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit:
|
||||
summary = image_studio_generation.run_jobs(
|
||||
[job],
|
||||
config=cfg,
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
|
||||
self.assertEqual(1, summary["failed"])
|
||||
self.assertIn("参考图尚未下载", summary["jobs"][0]["error"])
|
||||
submit.assert_not_called()
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_stop_after_download_discards_temporary_result(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, source = self._project_source(temp_dir)
|
||||
|
||||
@@ -99,6 +99,7 @@ class ProductSuiteTests(unittest.TestCase):
|
||||
|
||||
self.assertEqual(3, len(specs))
|
||||
self.assertEqual([11, 11, 12], [spec["source_asset_id"] for spec in specs])
|
||||
self.assertEqual([[], [], []], [spec["reference_asset_ids"] for spec in specs])
|
||||
self.assertEqual(["白底图", "场景图", "场景图"], [spec["job_type"] for spec in specs])
|
||||
for spec in specs:
|
||||
self.assertIn("平台:Shopee", spec["prompt"])
|
||||
@@ -135,6 +136,26 @@ class ProductSuiteTests(unittest.TestCase):
|
||||
self.assertLess(white_prompt.index("尺码信息规则"), white_prompt.index(reference_rule))
|
||||
self.assertLess(white_prompt.index(reference_rule), white_prompt.index("商品卖点与要求"))
|
||||
|
||||
def test_job_specs_freeze_first_image_references_when_not_per_image_primary(self):
|
||||
settings = product_suite.default_suite_settings()
|
||||
settings.update(
|
||||
{
|
||||
"per_image_primary": False,
|
||||
"categories": {"白底图": 1, "场景图": 1, "卖点图": 0},
|
||||
}
|
||||
)
|
||||
specs = product_suite.build_job_specs(
|
||||
[SimpleNamespace(id=11), SimpleNamespace(id=12), SimpleNamespace(id=13)],
|
||||
"卖点",
|
||||
settings,
|
||||
"51100639510",
|
||||
template_text=prompts.load_default_product_suite_prompt(),
|
||||
)
|
||||
|
||||
self.assertEqual(2, len(specs))
|
||||
self.assertEqual([11, 11], [spec["source_asset_id"] for spec in specs])
|
||||
self.assertEqual([[12, 13], [12, 13]], [spec["reference_asset_ids"] for spec in specs])
|
||||
|
||||
def test_reference_rule_follows_per_image_primary_setting(self):
|
||||
template = prompts.load_default_product_suite_prompt()
|
||||
settings = product_suite.default_suite_settings()
|
||||
|
||||
Reference in New Issue
Block a user