wallpaper.py: add randomly chosen perspective to model prompt, be explicit about device, add sampling knobs

1 files changed, 83 insertions, 45 deletions

diff --git a/wallpaper.py b/wallpaper.py
index c60b8c5..c4d9f79 100755
--- a/wallpaper.py
+++ b/wallpaper.py
@@ -28,11 +28,9 @@
 # url = "https://download.pytorch.org/whl/cpu"
 # explicit = true
 # ///
-
 import os
 
 OFFLINE = os.environ.get("OFFLINE", "1") != "0"
-
 os.environ["HF_HUB_OFFLINE"] = "1" if OFFLINE else "0"
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 
@@ -42,6 +40,7 @@ import argparse
 import datetime as dt
 import logging
 import random
+import re
 import subprocess
 import sys
 import gc
@@ -61,10 +60,9 @@ from spandrel import ModelLoader
 from suntimes import SunTimes
 
 logger = logging.getLogger("wallpaper")
-
-
 CAP = torch.backends.cpu.get_cpu_capability()
 COMPUTE_DTYPE = torch.bfloat16 if CAP.startswith("AVX512") else torch.float32
+DEVICE = "cpu"
 
 
 def step_context(
@@ -73,7 +71,6 @@ def step_context(
     longitude: float | None,
     elevation: int = 0,
 ) -> str:
-
     Y = 2000  # dummy leap year to allow input X-02-29 (leap day)
     seasons = [
         ("winter", (dt.date(Y, 1, 1), dt.date(Y, 3, 20))),
@@ -82,18 +79,15 @@ def step_context(
         ("autumn", (dt.date(Y, 9, 23), dt.date(Y, 12, 20))),
         ("winter", (dt.date(Y, 12, 21), dt.date(Y, 12, 31))),
     ]
-
     date = now.date()
     date = date.replace(year=Y)
     season = next(season for season, (start, end) in seasons if start <= date <= end)
-
     today = now.strftime(f"Date and time: %A %B %d, %Y at %H:%M %Z. Season: {season}.")
 
     if latitude is None or longitude is None:
         return today
 
     sun = SunTimes(longitude, latitude, elevation)
-
     today_sr = sun.riselocal(now)
     today_ss = sun.setlocal(now)
     sun_times = (
@@ -147,29 +141,36 @@ def step_llm(
     system_prompt: str,
     max_tokens: int,
     taste_prompts: list[str],
+    perspectives: list[str],
+    temperature: float,
+    top_p: float,
+    min_p: float,
     ctx: str,
 ) -> str:
     logger.info("Loading llm model...")
-
     model_path = hf_hub_download(
         repo_id=repo_id,
         filename=filename,
     )
-
     llm = Llama(
         model_path=model_path,
         n_ctx=2048,
         verbose=False,
     )
-    taste_prompt: list[str] = list(taste_prompts)
-    random.shuffle(taste_prompt)
-    taste_prompt_str = ", ".join(taste_prompt)
+
+    taste_prompts = list(taste_prompts)
+    random.shuffle(taste_prompts)
+    taste_prompt_str = ", ".join(taste_prompts)
+    perspective = random.choice(perspectives)
+    logger.info(f"Perspective: {perspective}")
 
     prompt = (
-        f"Taste reference:\n{taste_prompt_str}\n\n"
-        f"Context: {ctx}\n\n"
+        f"Taste reference: {taste_prompt_str}. "
+        f"Required perspective: {perspective}. "
+        f"Context: {ctx}. "
         f"Reply with one text-to-image prompt only."
     )
+
     logger.info("Running llm...")
     output = llm.create_chat_completion(
         messages=[
@@ -180,11 +181,18 @@ def step_llm(
             },
         ],
         max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        min_p=min_p,
     )
-
     del llm
     gc.collect()
-    return output["choices"][0]["message"]["content"].strip()
+
+    # might as well
+    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()
 
 
 def step_diffuse(
@@ -195,7 +203,6 @@ def step_diffuse(
     steps: int,
     prompt: str,
 ) -> list[Any]:
-
     ckpt_path = f"https://huggingface.co/{repo_id}/blob/main/{filename}"
 
     # populate cache
@@ -210,7 +217,7 @@ def step_diffuse(
         vae=None,
         torch_dtype=torch.bfloat16,  # always, for memory purposes
         local_files_only=OFFLINE,
-    )
+    ).to(DEVICE)
     with torch.inference_mode():
         prompt_embeds, _text_ids = pipe.encode_prompt(prompt=prompt)
     del pipe
@@ -227,7 +234,7 @@ def step_diffuse(
         subfolder="transformer",
         device_map="cpu",
         local_files_only=OFFLINE,
-    )
+    ).to(DEVICE)
     pipe = Flux2KleinPipeline.from_pretrained(
         "black-forest-labs/FLUX.2-klein-4B",
         text_encoder=None,
@@ -235,16 +242,17 @@ def step_diffuse(
         transformer=transformer,
         torch_dtype=COMPUTE_DTYPE,
         local_files_only=OFFLINE,
-    )
-
+    ).to(DEVICE)
     out = pipe(
         prompt_embeds=prompt_embeds,
         width=gen_width,
         height=gen_height,
         num_inference_steps=steps,
     )
+
     del pipe, transformer, prompt_embeds
     gc.collect()
+
     return out.images
 
 
@@ -289,35 +297,35 @@ def step_upscale(
 ) -> list[Any]:
     model_path = hf_hub_download(repo_id=repo_id, filename=filename)
     model = ModelLoader().load_from_file(model_path)
-    model.to("cpu").eval()
-
+    model.to(DEVICE).eval()
     results: list[Any] = []
     for im in images:
         arr = np.array(im.convert("RGB"), dtype=np.uint8)
-        x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).float().div(255)
-
+        x = (
+            torch.from_numpy(arr)
+            .permute(2, 0, 1)
+            .unsqueeze(0)
+            .float()
+            .div(255)
+            .to(DEVICE)
+        )
         if tile > 0:
             up = upscale_tiled(model, x, model.scale, tile=tile)
         else:
             with torch.inference_mode():
                 up = model(x)
-
         # Tensor is [1, C, H', W'] -> Permute to [H', W', C], scale from
         # [0, 1] to uint8 and clamp
         img_array = up[0].permute(1, 2, 0).mul(255).clamp(0, 255).byte().numpy()
-
         del x, up
         gc.collect()
         final_img = Image.fromarray(img_array)
-
         # by default 4x lands above the target (1280x720 -> 5120x2880), then
         # Lanczos downscales to the target, averaging away upscaler noise.
         target_size = (target_width, target_height)
         if (final_img.width, final_img.height) != target_size:
             final_img = final_img.resize(target_size, Image.LANCZOS)
-
         results.append(final_img)
-
     del model
     gc.collect()
     return results
@@ -334,7 +342,6 @@ def step_save(
         path = output_dir / fname
         img.save(path, "JPEG", quality=95)
         saved.append(path)
-
     txt_path = output_dir / f"prompt-{tag}.txt"
     txt_path.write_text(prompt_text)
     return saved
@@ -356,7 +363,6 @@ def parse_args() -> argparse.Namespace:
         description="Generate wallpaper. Use OFFLINE=0 env var for first run.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-
     wall = parser.add_argument_group("wallpaper")
     wall.add_argument(
         "--taste-prompt",
@@ -367,19 +373,36 @@ def parse_args() -> argparse.Namespace:
             "trees",
             "plants",
             "landscape",
+            "brutalist concrete",
             "hypermodern concrete architecture",
             "moss",
             "rocks",
             "abstract concrete sculptures",
             "precise water channels",
-            "futuristic megastructure of strange shape",
-            "a choice of perspective: perhaps internal, arial, distant or in-situ",
+            "megastructure of strange shape",
             "distinct scales of objects, some very large for contrast",
             "distant structures",
         ],
         help="Taste keywords for the LLM art-director prompt.",
     )
     wall.add_argument(
+        "--perspectives",
+        nargs="+",
+        default=[
+            "low ground-level view looking up",
+            "aerial view from great height",
+            "interior view looking outward through an opening",
+            "interior view looking inward to a courtyard",
+            "distant view across open landscape, structures on the horizon",
+            "in-situ view from within the structures, human eye level",
+            "view from across a body of water",
+            "view from a field",
+            "framed view through a narrow gap or corridor",
+        ],
+        help="Camera perspectives; one is chosen at random per run and "
+        "injected as a hard directive.",
+    )
+    wall.add_argument(
         "--ephemeral-dir", default="/tmp/wallpaper/", help="Output directory."
     )
     wall.add_argument(
@@ -399,7 +422,6 @@ def parse_args() -> argparse.Namespace:
         default=["awww", "img"],
         help="Command to run when setting wallpaper.",
     )
-
     llm = parser.add_argument_group("llm")
     llm.add_argument(
         "--llm-repo-id",
@@ -418,12 +440,31 @@ def parse_args() -> argparse.Namespace:
         help="Max tokens of prompt to generate.",
     )
     llm.add_argument(
+        "--llm-temperature",
+        type=float,
+        default=0.4,
+        help="Sampling temperature; llama-cpp-python's default is 0.2",
+    )
+    llm.add_argument(
+        "--llm-top-p",
+        type=float,
+        default=0.95,
+        help="Nucleus sampling threshold.",
+    )
+    llm.add_argument(
+        "--llm-min-p",
+        type=float,
+        default=0.05,
+        help="Minimum token probability relative to the top token",
+    )
+    llm.add_argument(
         "--llm-system-prompt",
         default=(
             "You are a wallpaper art director. "
             "Given the user's taste and the context, write one detailed text-to-image prompt for a diffusion model. "
             "You are repeatedly called upon to do this, and must introduce variation into your output. "
-            "Not all of the user's taste should be present, it's a suggestion and you should choose only a subset and improvise on the themes. "
+            "The taste reference lists the themes to build on; improvise around them, or drop some. "
+            "The required perspective MUST be used as the camera viewpoint of the image. "
             "Times after sunset and before sunrise should generate dark images for night time. "
             "Times during the day should track the brightness of the hour. "
             "The weather should be reflected in your prompt. "
@@ -433,7 +474,6 @@ def parse_args() -> argparse.Namespace:
         ),
         help="System prompt for LLM.",
     )
-
     flux2 = parser.add_argument_group("Flux Klein 4B settings")
     flux2.add_argument(
         "--flux2-klein-repo-id",
@@ -460,7 +500,6 @@ def parse_args() -> argparse.Namespace:
         default=720,
         help="Generated image height.",
     )
-
     sr = parser.add_argument_group("super-resolution")
     sr.add_argument(
         "--sr-repo-id",
@@ -480,7 +519,6 @@ def parse_args() -> argparse.Namespace:
         help="Tile size for tiled upscaling; bounds peak memory on "
         "constrained machines. 0 = full-frame (best quality).",
     )
-
     wx = parser.add_argument_group("weather")
     wx.add_argument(
         "--latitude", type=float, default=51.75, help="Latitude for weather."
@@ -494,7 +532,6 @@ def parse_args() -> argparse.Namespace:
         default=61,
         help="Altitude (m) for sunrise/sunset times.",
     )
-
     return parser.parse_args()
 
 
@@ -505,12 +542,10 @@ def run() -> None:
         datefmt="%H:%M:%S",
         stream=sys.stderr,
     )
-
     args = parse_args()
 
     # crazy
     now = dt.datetime.now(dt.datetime.now().astimezone().tzinfo)
-
     logger.info(f"Capability {CAP}, dtype {COMPUTE_DTYPE}")
 
     # 1. context
@@ -524,6 +559,10 @@ def run() -> None:
         args.llm_system_prompt,
         args.llm_max_tokens,
         list(args.taste_prompt),
+        list(args.perspectives),
+        args.llm_temperature,
+        args.llm_top_p,
+        args.llm_min_p,
         ctx,
     )
     logger.info(f"Prompt: {prompt_text}")
@@ -567,11 +606,10 @@ def run() -> None:
 
     # 5. save to configured output directory
     saved = step_save(upscaled, Path(args.ephemeral_dir).expanduser(), now, prompt_text)
-
     for p in saved:
         logger.info("Saved: %s", p.resolve())
-
     # 6. set wallpaper
+
     step_wallpaper(
         args.set_command,
         saved,