Restore clothing details for atlas prompts

This commit is contained in:
2026-07-01 00:06:27 +02:00
parent f8f2fb43df
commit 7e41613c1e
2 changed files with 79 additions and 0 deletions
+24
View File
@@ -145,6 +145,26 @@ def _restore_axis_values_for_context(
return values return values
def _fallback_restore_axis_values(
source_categories: list[dict[str, Any]] | None,
axis_name: str,
hardcore_position_config: dict[str, Any],
) -> list[Any]:
if not source_categories:
return []
values: list[Any] = []
for category in source_categories:
if not hardcore_position_policy.is_hardcore_sexual_category(category):
continue
for subcategory in category.get("subcategories", []):
raw_axes = subcategory.get("item_axes")
if isinstance(raw_axes, dict):
values.extend(_list_from(raw_axes.get(axis_name)))
if not values:
return []
return hardcore_position_policy.filter_hardcore_axis(axis_name, values, hardcore_position_config)
def _restored_prompt_axis_values( def _restored_prompt_axis_values(
rng: Any, rng: Any,
subcategory: dict[str, Any], subcategory: dict[str, Any],
@@ -152,6 +172,7 @@ def _restored_prompt_axis_values(
hardcore_position_config: dict[str, Any], hardcore_position_config: dict[str, Any],
women_count: int, women_count: int,
men_count: int, men_count: int,
source_categories: list[dict[str, Any]] | None = None,
) -> dict[str, str]: ) -> dict[str, str]:
restore_axes = hardcore_position_policy.normalize_restore_prompt_axes( restore_axes = hardcore_position_policy.normalize_restore_prompt_axes(
hardcore_position_config.get("restore_prompt_axes") if isinstance(hardcore_position_config, dict) else [] hardcore_position_config.get("restore_prompt_axes") if isinstance(hardcore_position_config, dict) else []
@@ -170,6 +191,8 @@ def _restored_prompt_axis_values(
restored[axis_name] = existing restored[axis_name] = existing
continue continue
values = _list_from(raw_axes.get(axis_name)) values = _list_from(raw_axes.get(axis_name))
if not values:
values = _fallback_restore_axis_values(source_categories, axis_name, hardcore_position_config)
if not values: if not values:
continue continue
values = _restore_axis_values_for_context( values = _restore_axis_values_for_context(
@@ -261,6 +284,7 @@ def select_category_item_route_result(
parsed_hardcore_position_config, parsed_hardcore_position_config,
women_count, women_count,
men_count, men_count,
source_categories,
) )
if restored_axis_values: if restored_axis_values:
restored_details = _unique_texts(list(restored_axis_values.values())) restored_details = _unique_texts(list(restored_axis_values.values()))
+55
View File
@@ -8519,6 +8519,61 @@ def smoke_pov_oral_position_routes() -> None:
any(str(detail).lower() in restored_top_prompt for detail in restored_details), any(str(detail).lower() in restored_top_prompt for detail in restored_details),
f"Krea2 POV Prompt Restore final prompt did not include sampled restored details: {restored_top_prompt}", f"Krea2 POV Prompt Restore final prompt did not include sampled restored details: {restored_top_prompt}",
) )
clothing_only_top_config = sxcp_nodes.NODE_CLASS_MAPPINGS["SxCPKrea2POVPromptRestore"]().build(
True,
False,
False,
False,
True,
top_variant_config,
)[0]
clothing_only_top_pair = pb.build_insta_of_pair(
row_number=1,
start_index=1,
seed=3828,
ethnicity="any",
figure="random",
no_plus_women=False,
no_black=False,
trigger=Trigger,
prepend_trigger_to_prompt=True,
options_json=_insta_options(
softcore_camera_mode="from_camera_config",
hardcore_camera_mode="from_camera_config",
camera_detail="compact",
),
character_cast=_character_cast(pov_man=True),
hardcore_position_config=clothing_only_top_config,
location_config=_coworking_location_config(),
hardcore_camera_config=_orbit_camera(
horizontal_angle=45,
vertical_angle=0,
zoom=7.5,
subject_focus="action",
),
)
clothing_only_top_row = clothing_only_top_pair.get("hardcore_row") or {}
clothing_only_top_axis = clothing_only_top_row.get("item_axis_values") or {}
clothing_only_axes = clothing_only_top_axis.get("restored_prompt_axes") or []
clothing_only_details = clothing_only_top_axis.get("restored_prompt_details") or []
_expect(
clothing_only_axes == ["clothing_detail"],
f"Krea2 POV Prompt Restore clothing-only mode should restore clothing_detail, got {clothing_only_axes!r}",
)
_expect(
clothing_only_details,
"Krea2 POV Prompt Restore clothing-only mode should add a sampled clothing detail",
)
clothing_only_top_krea = krea_formatter.format_krea2_prompt("", metadata_json=_json(clothing_only_top_pair), target="hardcore")
clothing_only_top_prompt = _expect_text(
"pov_oral_top_view_restore_clothing_only.krea_prompt",
clothing_only_top_krea.get("krea_prompt"),
60,
).lower()
_expect(
any(str(detail).lower() in clothing_only_top_prompt for detail in clothing_only_details),
f"Krea2 POV Prompt Restore clothing-only final prompt did not include sampled clothing detail: {clothing_only_top_prompt}",
)
sitting_variant_config = sxcp_nodes.NODE_CLASS_MAPPINGS["SxCPKrea2POVOralFilter"]().build( sitting_variant_config = sxcp_nodes.NODE_CLASS_MAPPINGS["SxCPKrea2POVOralFilter"]().build(
"replace", "replace",