Restore clothing details for atlas prompts
This commit is contained in:
@@ -145,6 +145,26 @@ def _restore_axis_values_for_context(
|
||||
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(
|
||||
rng: Any,
|
||||
subcategory: dict[str, Any],
|
||||
@@ -152,6 +172,7 @@ def _restored_prompt_axis_values(
|
||||
hardcore_position_config: dict[str, Any],
|
||||
women_count: int,
|
||||
men_count: int,
|
||||
source_categories: list[dict[str, Any]] | None = None,
|
||||
) -> dict[str, str]:
|
||||
restore_axes = hardcore_position_policy.normalize_restore_prompt_axes(
|
||||
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
|
||||
continue
|
||||
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:
|
||||
continue
|
||||
values = _restore_axis_values_for_context(
|
||||
@@ -261,6 +284,7 @@ def select_category_item_route_result(
|
||||
parsed_hardcore_position_config,
|
||||
women_count,
|
||||
men_count,
|
||||
source_categories,
|
||||
)
|
||||
if restored_axis_values:
|
||||
restored_details = _unique_texts(list(restored_axis_values.values()))
|
||||
|
||||
@@ -8519,6 +8519,61 @@ def smoke_pov_oral_position_routes() -> None:
|
||||
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}",
|
||||
)
|
||||
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(
|
||||
"replace",
|
||||
|
||||
Reference in New Issue
Block a user