wallpaper.py: explicit instructions to flux

1 files changed, 10 insertions, 2 deletions

diff --git a/wallpaper.py b/wallpaper.py
index 39e8929..07a7c6e 100755
--- a/wallpaper.py
+++ b/wallpaper.py
@@ -148,6 +148,7 @@ def step_llm(
     top_p: float,
     min_p: float,
     ctx: str,
+    explicit_instr: str,
 ) -> str:
     logger.info("Loading llm model...")
     model_path = hf_hub_download(
@@ -200,7 +201,8 @@ def step_llm(
     text = output["choices"][0]["message"]["content"]
     text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
     text = re.sub(r"<think>.*", "", text, flags=re.DOTALL)
-    return text.strip()
+    text = text.strip()
+    return text + " " + explicit_instr
 
 
 def step_diffuse(
@@ -478,7 +480,7 @@ def parse_args() -> argparse.Namespace:
     llm.add_argument(
         "--llm-max-tokens",
         type=int,
-        default=200,
+        default=384,
         help="Max tokens of prompt to generate.",
     )
     llm.add_argument(
@@ -521,6 +523,11 @@ def parse_args() -> argparse.Namespace:
     )
     flux2 = parser.add_argument_group("Flux Klein 4B settings")
     flux2.add_argument(
+        "--flux2-explicit-instructions",
+        default="No railings, no glass, no people, no objects.",
+        help="Explicit instructions always provided to the diffusion model.",
+    )
+    flux2.add_argument(
         "--flux2-klein-repo-id",
         default="unsloth/FLUX.2-klein-4B-GGUF",
         help="HuggingFace repo for quantised model.",
@@ -609,6 +616,7 @@ def run() -> None:
         args.llm_top_p,
         args.llm_min_p,
         ctx,
+        args.flux2_explicit_instructions,
     )
     logger.info(f"Prompt: {prompt_text}")