Spaces:

build-small-hackathon
/

compliment-forest

Running

compliment-forest / src /compliment_forest /style_data.py

fix: synchronize the complete application runtime (#5)

513ea7b 5 months ago

7.6 kB

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

from .backends.image import STYLE_PROFILES, compose_flux_prompt

SceneCategory = Literal["animal", "human", "object", "environment"]


@dataclass(frozen=True)
class ForestScene:
    slug: str
    category: SceneCategory
    prompt: str


FOREST_SCENES = (
    ForestScene(
        "fox-threshold",
        "animal",
        "a gentle red fox pausing at the edge of a fern-lined path",
    ),
    ForestScene(
        "listening-owl",
        "animal",
        "a round barn owl listening from a low mossy branch",
    ),
    ForestScene(
        "steady-deer",
        "animal",
        "a young deer standing calmly between silver birch trees",
    ),
    ForestScene(
        "brave-snail",
        "animal",
        "a tiny snail crossing a dew-covered fern at dawn",
    ),
    ForestScene(
        "singing-wren",
        "animal",
        "a small wren singing beside loose woodland flowers",
    ),
    ForestScene(
        "river-otter",
        "animal",
        "a river otter holding one smooth stone beside quiet reeds",
    ),
    ForestScene(
        "thoughtful-badger",
        "animal",
        "a thoughtful badger beside a lantern-shaped mushroom",
    ),
    ForestScene(
        "patient-hare",
        "animal",
        "a patient brown hare resting beneath arching grasses",
    ),
    ForestScene(
        "moonlit-moth",
        "animal",
        "a luna moth hovering near moonlit foxgloves",
    ),
    ForestScene(
        "walking-turtle",
        "animal",
        "a small woodland turtle moving between clover and stones",
    ),
    ForestScene(
        "person-open-window",
        "human",
        "an adult seen from behind opening a window to pale morning light",
    ),
    ForestScene(
        "person-blank-notebook",
        "human",
        "an adult seated at a wooden desk with an open blank notebook",
    ),
    ForestScene(
        "person-forked-path",
        "human",
        "a small human figure viewed from behind at a gentle fork in a path",
    ),
    ForestScene(
        "person-train-platform",
        "human",
        "a quiet adult figure waiting on a misty train platform with one bag",
    ),
    ForestScene(
        "person-moving-box",
        "human",
        "an adult carrying one moving box toward a sunlit doorway",
    ),
    ForestScene(
        "person-footbridge",
        "human",
        "a side-view figure taking one step across a narrow wooden footbridge",
    ),
    ForestScene(
        "person-doorway",
        "human",
        "a calm adult silhouette standing in an open doorway between two rooms",
    ),
    ForestScene(
        "person-seedling",
        "human",
        "hands gently watering a small seedling on a windowsill",
    ),
    ForestScene(
        "person-rain-shelter",
        "human",
        "an adult seen from the side resting on a bench beneath a rain shelter",
    ),
    ForestScene(
        "person-dawn-hill",
        "human",
        "a distant human figure standing on a low hillside at dawn",
    ),
    ForestScene(
        "lantern-crossroads",
        "object",
        "a small glowing lantern placed where two woodland paths meet",
    ),
    ForestScene(
        "map-compass",
        "object",
        "an unfolded map and simple compass resting on a wooden table",
    ),
    ForestScene(
        "open-notebook",
        "object",
        "an open blank notebook beside a pencil and one pressed leaf",
    ),
    ForestScene(
        "stepping-stones",
        "object",
        "four smooth stepping stones crossing a narrow stream",
    ),
    ForestScene(
        "warm-cup",
        "object",
        "a warm ceramic cup sending a thin curl of steam into morning light",
    ),
    ForestScene(
        "woven-thread",
        "object",
        "loose green and gold threads gradually woven into one calm pattern",
    ),
    ForestScene(
        "key-and-door",
        "object",
        "a simple brass key resting beside a small unopened wooden door",
    ),
    ForestScene(
        "paper-boat",
        "object",
        "a single paper boat floating on still water beneath willow reflections",
    ),
    ForestScene(
        "balanced-stones",
        "object",
        "three imperfect river stones balanced beside soft grasses",
    ),
    ForestScene(
        "empty-chair-light",
        "object",
        "an empty wooden chair in a quiet patch of warm window light",
    ),
    ForestScene(
        "winding-path",
        "environment",
        "a winding path disappearing gently through tall ferns and morning mist",
    ),
    ForestScene(
        "river-crossing",
        "environment",
        "a shallow river crossing with stones visible beneath clear water",
    ),
    ForestScene(
        "room-at-dawn",
        "environment",
        "a quiet room at dawn with curtains moving beside an open window",
    ),
    ForestScene(
        "city-garden",
        "environment",
        "a small green garden between quiet city buildings after rain",
    ),
    ForestScene(
        "misty-platform",
        "environment",
        "an empty train platform fading softly into early morning mist",
    ),
    ForestScene(
        "clearing-after-rain",
        "environment",
        "a forest clearing just after rain with one bright opening in the clouds",
    ),
    ForestScene(
        "hillside-trail",
        "environment",
        "a gradual hillside trail curving toward a pale open horizon",
    ),
    ForestScene(
        "staircase-light",
        "environment",
        "a simple staircase with warm light falling across the next three steps",
    ),
    ForestScene(
        "canopy-opening",
        "environment",
        "a dark green canopy opening into a circle of soft sky",
    ),
    ForestScene(
        "shoreline-horizon",
        "environment",
        "a calm shoreline where fading clouds meet a wide quiet horizon",
    ),
)

# Compatibility alias for callers that used the v1 name.
FOREST_SUBJECTS = FOREST_SCENES

TRAINED_STYLE_IDS = (
    "watercolor",
    "paper_cut",
    "moonlit_gouache",
    "botanical_ink",
)


def build_style_records(
    *,
    samples_per_style: int = 40,
    base_seed: int = 9000,
) -> list[dict[str, str | int]]:
    if not 1 <= samples_per_style <= len(FOREST_SCENES):
        raise ValueError(f"samples_per_style must be between 1 and {len(FOREST_SCENES)}")

    records: list[dict[str, str | int]] = []
    for style_offset, style in enumerate(TRAINED_STYLE_IDS):
        profile = STYLE_PROFILES[style]
        for scene_index, scene in enumerate(FOREST_SCENES[:samples_per_style]):
            seed = base_seed + style_offset * 1000 + scene_index
            prompt = compose_flux_prompt(
                scene.prompt,
                style=style,  # type: ignore[arg-type]
                seed=seed,
            )
            records.append(
                {
                    "style": style,
                    "trigger": profile.trigger,
                    "category": scene.category,
                    "subject": scene.slug,
                    "seed": seed,
                    "prompt": prompt,
                    "text": (
                        f"{profile.trigger}, {scene.prompt}, "
                        f"{profile.label.lower()} storybook scene"
                    ),
                    "file_name": (
                        f"{scene_index:03d}-{scene.category}-{scene.slug}.png"
                    ),
                }
            )
    return records

	from __future__ import annotations

	from dataclasses import dataclass
	from typing import Literal

	from .backends.image import STYLE_PROFILES, compose_flux_prompt

	SceneCategory = Literal["animal", "human", "object", "environment"]


	@dataclass(frozen=True)
	class ForestScene:
	slug: str
	category: SceneCategory
	prompt: str


	FOREST_SCENES = (
	ForestScene(
	"fox-threshold",
	"animal",
	"a gentle red fox pausing at the edge of a fern-lined path",
	),
	ForestScene(
	"listening-owl",
	"animal",
	"a round barn owl listening from a low mossy branch",
	),
	ForestScene(
	"steady-deer",
	"animal",
	"a young deer standing calmly between silver birch trees",
	),
	ForestScene(
	"brave-snail",
	"animal",
	"a tiny snail crossing a dew-covered fern at dawn",
	),
	ForestScene(
	"singing-wren",
	"animal",
	"a small wren singing beside loose woodland flowers",
	),
	ForestScene(
	"river-otter",
	"animal",
	"a river otter holding one smooth stone beside quiet reeds",
	),
	ForestScene(
	"thoughtful-badger",
	"animal",
	"a thoughtful badger beside a lantern-shaped mushroom",
	),
	ForestScene(
	"patient-hare",
	"animal",
	"a patient brown hare resting beneath arching grasses",
	),
	ForestScene(
	"moonlit-moth",
	"animal",
	"a luna moth hovering near moonlit foxgloves",
	),
	ForestScene(
	"walking-turtle",
	"animal",
	"a small woodland turtle moving between clover and stones",
	),
	ForestScene(
	"person-open-window",
	"human",
	"an adult seen from behind opening a window to pale morning light",
	),
	ForestScene(
	"person-blank-notebook",
	"human",
	"an adult seated at a wooden desk with an open blank notebook",
	),
	ForestScene(
	"person-forked-path",
	"human",
	"a small human figure viewed from behind at a gentle fork in a path",
	),
	ForestScene(
	"person-train-platform",
	"human",
	"a quiet adult figure waiting on a misty train platform with one bag",
	),
	ForestScene(
	"person-moving-box",
	"human",
	"an adult carrying one moving box toward a sunlit doorway",
	),
	ForestScene(
	"person-footbridge",
	"human",
	"a side-view figure taking one step across a narrow wooden footbridge",
	),
	ForestScene(
	"person-doorway",
	"human",
	"a calm adult silhouette standing in an open doorway between two rooms",
	),
	ForestScene(
	"person-seedling",
	"human",
	"hands gently watering a small seedling on a windowsill",
	),
	ForestScene(
	"person-rain-shelter",
	"human",
	"an adult seen from the side resting on a bench beneath a rain shelter",
	),
	ForestScene(
	"person-dawn-hill",
	"human",
	"a distant human figure standing on a low hillside at dawn",
	),
	ForestScene(
	"lantern-crossroads",
	"object",
	"a small glowing lantern placed where two woodland paths meet",
	),
	ForestScene(
	"map-compass",
	"object",
	"an unfolded map and simple compass resting on a wooden table",
	),
	ForestScene(
	"open-notebook",
	"object",
	"an open blank notebook beside a pencil and one pressed leaf",
	),
	ForestScene(
	"stepping-stones",
	"object",
	"four smooth stepping stones crossing a narrow stream",
	),
	ForestScene(
	"warm-cup",
	"object",
	"a warm ceramic cup sending a thin curl of steam into morning light",
	),
	ForestScene(
	"woven-thread",
	"object",
	"loose green and gold threads gradually woven into one calm pattern",
	),
	ForestScene(
	"key-and-door",
	"object",
	"a simple brass key resting beside a small unopened wooden door",
	),
	ForestScene(
	"paper-boat",
	"object",
	"a single paper boat floating on still water beneath willow reflections",
	),
	ForestScene(
	"balanced-stones",
	"object",
	"three imperfect river stones balanced beside soft grasses",
	),
	ForestScene(
	"empty-chair-light",
	"object",
	"an empty wooden chair in a quiet patch of warm window light",
	),
	ForestScene(
	"winding-path",
	"environment",
	"a winding path disappearing gently through tall ferns and morning mist",
	),
	ForestScene(
	"river-crossing",
	"environment",
	"a shallow river crossing with stones visible beneath clear water",
	),
	ForestScene(
	"room-at-dawn",
	"environment",
	"a quiet room at dawn with curtains moving beside an open window",
	),
	ForestScene(
	"city-garden",
	"environment",
	"a small green garden between quiet city buildings after rain",
	),
	ForestScene(
	"misty-platform",
	"environment",
	"an empty train platform fading softly into early morning mist",
	),
	ForestScene(
	"clearing-after-rain",
	"environment",
	"a forest clearing just after rain with one bright opening in the clouds",
	),
	ForestScene(
	"hillside-trail",
	"environment",
	"a gradual hillside trail curving toward a pale open horizon",
	),
	ForestScene(
	"staircase-light",
	"environment",
	"a simple staircase with warm light falling across the next three steps",
	),
	ForestScene(
	"canopy-opening",
	"environment",
	"a dark green canopy opening into a circle of soft sky",
	),
	ForestScene(
	"shoreline-horizon",
	"environment",
	"a calm shoreline where fading clouds meet a wide quiet horizon",
	),
	)

	# Compatibility alias for callers that used the v1 name.
	FOREST_SUBJECTS = FOREST_SCENES

	TRAINED_STYLE_IDS = (
	"watercolor",
	"paper_cut",
	"moonlit_gouache",
	"botanical_ink",
	)


	def build_style_records(
	*,
	samples_per_style: int = 40,
	base_seed: int = 9000,
	) -> list[dict[str, str \| int]]:
	if not 1 <= samples_per_style <= len(FOREST_SCENES):
	raise ValueError(f"samples_per_style must be between 1 and {len(FOREST_SCENES)}")

	records: list[dict[str, str \| int]] = []
	for style_offset, style in enumerate(TRAINED_STYLE_IDS):
	profile = STYLE_PROFILES[style]
	for scene_index, scene in enumerate(FOREST_SCENES[:samples_per_style]):
	seed = base_seed + style_offset * 1000 + scene_index
	prompt = compose_flux_prompt(
	scene.prompt,
	style=style, # type: ignore[arg-type]
	seed=seed,
	)
	records.append(
	{
	"style": style,
	"trigger": profile.trigger,
	"category": scene.category,
	"subject": scene.slug,
	"seed": seed,
	"prompt": prompt,
	"text": (
	f"{profile.trigger}, {scene.prompt}, "
	f"{profile.label.lower()} storybook scene"
	),
	"file_name": (
	f"{scene_index:03d}-{scene.category}-{scene.slug}.png"
	),
	}
	)
	return records