diff options
Diffstat (limited to 'wallpaper.py')
| -rwxr-xr-x | wallpaper.py | 905 |
1 files changed, 449 insertions, 456 deletions
diff --git a/wallpaper.py b/wallpaper.py index 07a7c6e..c98715b 100755 --- a/wallpaper.py +++ b/wallpaper.py @@ -3,66 +3,88 @@ # /// script # requires-python = ">=3.14" # dependencies = [ -# "accelerate>=1.14.0", -# "diffusers>=0.39.0", -# "gguf>=0.19.0", -# "llama-cpp-python>=0.3.34", -# "numpy>=2.5.1", -# "pillow>=12.3.0", # "requests>=2.34.2", -# "spandrel>=0.4.1", # "suntimes>=1.1.2", -# "torch>=2.13.0", -# "torchvision>=0.28.0", -# "transformers>=5.14.1", # ] -# [tool.uv.sources] -# torch = [ -# { index = "pytorch-cpu" }, -# ] -# torchvision = [ -# { index = "pytorch-cpu" }, -# ] -# [[tool.uv.index]] -# name = "pytorch-cpu" -# 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" -from pathlib import Path -from typing import Any import argparse import datetime as dt import logging import random import re +import shutil import subprocess import sys -import gc +from pathlib import Path -import numpy as np import requests -import torch -from llama_cpp import Llama -from huggingface_hub import hf_hub_download, snapshot_download -from diffusers import ( - Flux2KleinPipeline, - Flux2Transformer2DModel, - GGUFQuantizationConfig, -) -from PIL import Image -from spandrel import ModelLoader +from PIL import Image, ImageOps 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" + + +WMO_SHORT = { + 0: "clear sky", + 1: "mostly clear", + 2: "partly cloudy", + 3: "overcast", + 45: "foggy", + 48: "rime fog", + 51: "light drizzle", + 53: "drizzle", + 55: "heavy drizzle", + 61: "light rain", + 63: "rain", + 65: "heavy rain", + 71: "light snow", + 73: "snow", + 75: "heavy snow", + 80: "showers", + 81: "moderate showers", + 82: "heavy showers", + 95: "thunderstorm", + 96: "thunderstorm + hail", + 99: "severe thunderstorm", +} + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be greater than zero") + return parsed + + +def non_negative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be zero or greater") + return parsed + + +def non_negative_float(value: str) -> float: + parsed = float(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be zero or greater") + return parsed + + +def probability(value: str) -> float: + parsed = float(value) + if not 0.0 <= parsed <= 1.0: + raise argparse.ArgumentTypeError("must be between zero and one") + return parsed + + +def positive_fraction(value: str) -> float: + parsed = float(value) + if not 0.0 < parsed <= 1.0: + raise argparse.ArgumentTypeError( + "must be greater than zero and no greater than one" + ) + return parsed def step_context( @@ -71,74 +93,81 @@ def step_context( longitude: float | None, elevation: int = 0, ) -> str: - Y = 2000 # dummy leap year to allow input X-02-29 (leap day) + year = 2000 # Dummy leap year to allow February 29. seasons = [ - ("winter", (dt.date(Y, 1, 1), dt.date(Y, 3, 20))), - ("spring", (dt.date(Y, 3, 21), dt.date(Y, 6, 20))), - ("summer", (dt.date(Y, 6, 21), dt.date(Y, 9, 22))), - ("autumn", (dt.date(Y, 9, 23), dt.date(Y, 12, 20))), - ("winter", (dt.date(Y, 12, 21), dt.date(Y, 12, 31))), + ("winter", (dt.date(year, 1, 1), dt.date(year, 3, 20))), + ("spring", (dt.date(year, 3, 21), dt.date(year, 6, 20))), + ("summer", (dt.date(year, 6, 21), dt.date(year, 9, 22))), + ("autumn", (dt.date(year, 9, 23), dt.date(year, 12, 20))), + ("winter", (dt.date(year, 12, 21), dt.date(year, 12, 31))), ] - date = now.date() - date = date.replace(year=Y) - season = next(season for season, (start, end) in seasons if start <= date <= end) + + date = now.date().replace(year=year) + season = next( + season_name for season_name, (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) + sunrise = sun.riselocal(now) + sunset = sun.setlocal(now) sun_times = ( - f"sunrise {today_sr.strftime('%H:%M')}, sunset {today_ss.strftime('%H:%M')}" - ) - - detail = "unknown" - _WMO_SHORT = { - 0: "clear sky", - 1: "mostly clear", - 2: "partly cloudy", - 3: "overcast", - 45: "foggy", - 48: "rime fog", - 51: "light drizzle", - 53: "drizzle", - 55: "heavy drizzle", - 61: "light rain", - 63: "rain", - 65: "heavy rain", - 71: "light snow", - 73: "snow", - 75: "heavy snow", - 80: "showers", - 81: "moderate showers", - 82: "heavy showers", - 95: "thunderstorm", - 96: "thunderstorm + hail", - 99: "severe thunderstorm", - } + f"sunrise {sunrise.strftime('%H:%M')}, sunset {sunset.strftime('%H:%M')}" + ) + + weather = "unknown" try: - url = ( - f"https://api.open-meteo.com/v1/forecast?" - f"latitude={latitude}&longitude={longitude}" - f"¤t=weather_code,temperature_2m&timezone=auto" + response = requests.get( + "https://api.open-meteo.com/v1/forecast", + params={ + "latitude": latitude, + "longitude": longitude, + "current": "weather_code,temperature_2m", + "timezone": "auto", + }, + timeout=10, ) - data = requests.get(url, timeout=10).json()["current"] - code = data["weather_code"] - code_desc = _WMO_SHORT.get(code, "") - temp = float(data["temperature_2m"]) - detail = f"{code_desc} {temp:.0f}°C" - except Exception as exc: + response.raise_for_status() + + data = response.json()["current"] + code = int(data["weather_code"]) + description = WMO_SHORT.get(code, f"weather code {code}") + temperature = float(data["temperature_2m"]) + weather = f"{description}, {temperature:.0f}°C" + except (requests.RequestException, KeyError, TypeError, ValueError) as exc: logger.warning("Weather fetch failed (%s), continuing anyway.", exc) - return f"{today} Daylight: {sun_times}. Weather: {detail}" if detail else today + return f"{today} Daylight: {sun_times}. Weather: {weather}." + + +def clean_llm_output(text: str) -> str: + text = re.sub( + r"\s*\[\s*end of text\s*\]\s*$", + "", + text, + flags=re.IGNORECASE, + ) + text = re.sub( + r"<think>.*?</think>", + "", + text, + flags=re.DOTALL | re.IGNORECASE, + ) + text = re.sub( + r"<think>.*$", + "", + text, + flags=re.DOTALL | re.IGNORECASE, + ) + return text.strip() def step_llm( - repo_id: str, - filename: str, system_prompt: str, + model_path: Path, max_tokens: int, taste_prompts: list[str], taste_sample_frac: float, @@ -147,234 +176,187 @@ def step_llm( temperature: float, top_p: float, min_p: float, - ctx: str, - explicit_instr: str, + context: str, + explicit_instructions: 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, - ) + if not taste_prompts: + raise ValueError("At least one taste prompt is required") + if not biomes: + raise ValueError("At least one biome is required") + if not perspectives: + raise ValueError("At least one perspective is required") - # random.sample also shuffles, so ordering varies run to run total = len(taste_prompts) - k = max(1, round(taste_sample_frac * total)) - taste_prompts = random.sample(list(taste_prompts), k) - taste_prompt_str = ", ".join(taste_prompts) - logger.info(f"Taste ({k}/{total} sampled): {taste_prompt_str}") + sample_size = max(1, round(taste_sample_frac * total)) + sampled_tastes = random.sample(taste_prompts, sample_size) + taste_prompt_text = ", ".join(sampled_tastes) + biome = random.choice(biomes) - logger.info(f"Biome: {biome}") perspective = random.choice(perspectives) - logger.info(f"Perspective: {perspective}") - prompt = ( - f"Taste reference: {taste_prompt_str}. " + logger.info( + "Taste (%d/%d sampled): %s", + sample_size, + total, + taste_prompt_text, + ) + logger.info("Biome: %s", biome) + logger.info("Perspective: %s", perspective) + + user_prompt = ( + f"Taste reference: {taste_prompt_text}. " f"Required biome: {biome}. " f"Required perspective: {perspective}. " - f"Context: {ctx}. " - f"Reply with one text-to-image prompt only." - ) + f"Context: {context}. " + "Reply with one text-to-image prompt only. /nothink" + ) + + cmd = [ + "llama-completion", + "-m", + str(model_path.expanduser()), + "-c", + "2048", + "-n", + str(max_tokens), + "--temp", + str(temperature), + "--top-p", + str(top_p), + "--min-p", + str(min_p), + "--jinja", + "--single-turn", + "-sys", + system_prompt, + "-p", + user_prompt, + "--simple-io", + "--no-display-prompt", + "--no-perf", + "--log-verbosity", + "1", + ] - logger.info("Running llm...") - output = llm.create_chat_completion( - messages=[ - {"role": "system", "content": system_prompt}, - { - "role": "user", - "content": prompt, - }, - ], - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - min_p=min_p, - ) - del llm - gc.collect() + logger.info("Running llama.cpp ...") + try: + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + except FileNotFoundError as exc: + raise RuntimeError("Could not find the llama-completion executable") from exc + + if result.returncode != 0: + raise RuntimeError( + f"llama-completion exited with {result.returncode}:\n" + f"{result.stderr.strip()}" + ) + + generated_prompt = clean_llm_output(result.stdout) + if not generated_prompt: + raise RuntimeError( + "llama-completion completed successfully but produced no prompt" + ) - # 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) - text = text.strip() - return text + " " + explicit_instr + explicit_instructions = explicit_instructions.strip() + if explicit_instructions: + return f"{generated_prompt} {explicit_instructions}" + return generated_prompt def step_diffuse( - repo_id: str, - filename: str, - gen_width: int, - gen_height: int, + *, + diffusion_model: Path, + vae: Path, + llm: Path, + cfg_scale: float, steps: int, - prompt: str, -) -> list[Any]: - ckpt_path = f"https://huggingface.co/{repo_id}/blob/main/{filename}" - - # populate cache - if not OFFLINE: - snapshot_download("black-forest-labs/FLUX.2-klein-4B") - - # Phase 1: text encoding only - logger.info("Encoding prompt...") - pipe = Flux2KleinPipeline.from_pretrained( - "black-forest-labs/FLUX.2-klein-4B", - transformer=None, - 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 - gc.collect() - prompt_embeds = prompt_embeds.to(COMPUTE_DTYPE) - - # Phase 2: denoise + decode, without added memory load - logger.info("Denoising...") - transformer = Flux2Transformer2DModel.from_single_file( - ckpt_path, - quantization_config=GGUFQuantizationConfig(compute_dtype=COMPUTE_DTYPE), - torch_dtype=COMPUTE_DTYPE, - config="black-forest-labs/FLUX.2-klein-4B", - 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, - tokenizer=None, - 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 - - -def upscale_tiled( - model: Any, - x: torch.Tensor, - scale: int, - tile: int = 256, - overlap: int = 16, -) -> torch.Tensor: - # x: [1, C, H, W] in [0, 1]. Peak memory is bounded by the tile size, - # not the image size. Each tile is padded with `overlap` px of context, - # which is cropped away on write-back, so tiles blend seamlessly. - _, c, h, w = x.shape - out = torch.zeros(1, c, h * scale, w * scale) - stride = tile - 2 * overlap - for y0 in range(0, h, stride): - for x0 in range(0, w, stride): - iy0, ix0 = max(y0 - overlap, 0), max(x0 - overlap, 0) - iy1 = min(y0 + stride + overlap, h) - ix1 = min(x0 + stride + overlap, w) - with torch.inference_mode(): - up = model(x[:, :, iy0:iy1, ix0:ix1]) - oy0, ox0 = y0 * scale, x0 * scale - oy1 = min(y0 + stride, h) * scale - ox1 = min(x0 + stride, w) * scale - py0, px0 = (y0 - iy0) * scale, (x0 - ix0) * scale - out[:, :, oy0:oy1, ox0:ox1] = up[ - :, :, py0 : py0 + (oy1 - oy0), px0 : px0 + (ox1 - ox0) - ] - del up - return out - - -def step_upscale( - images: list[Any], - target_width: int, - target_height: int, - repo_id: str, - filename: str, - tile: int, -) -> list[Any]: - model_path = hf_hub_download(repo_id=repo_id, filename=filename) - model = ModelLoader().load_from_file(model_path) - 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) - .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 - - -def step_save( - upscaled: list[Any], - raw_images: list[Any], + flash_attention: bool, + generation_width: int, + generation_height: int, + upscale_model: Path, + upscale_repeats: int, + upscale_tile_size: int, output_dir: Path, now: dt.datetime, prompt_text: str, -) -> list[Path]: +) -> Path: output_dir.mkdir(parents=True, exist_ok=True) - saved = [] tag = now.strftime("%Y%m%d-%H%M%S") - for i, (up, raw) in enumerate(zip(upscaled, raw_images)): - idx = f"-{i}" if len(upscaled) > 1 else "" - for img, fname in ( - (up, f"wallpaper-{tag}{idx}.jpg"), - (raw, f"raw-{tag}{idx}.jpg"), - ): - path = output_dir / fname - img.save(path, "JPEG", quality=100) - if fname.startswith("wall"): - saved.append(path) - txt_path = output_dir / f"prompt-{tag}.txt" - txt_path.write_text(prompt_text) - return saved - - -def step_wallpaper(command: list[str], saved: list[Path], store: Path) -> None: - idx = random.randrange(len(saved)) + + output_path = output_dir / f"wallpaper-{tag}.png" + prompt_path = output_dir / f"prompt-{tag}.txt" + + cmd = [ + "sd-cli", + "--diffusion-model", + str(diffusion_model.expanduser()), + "--vae", + str(vae.expanduser()), + "--llm", + str(llm.expanduser()), + "--cfg-scale", + str(cfg_scale), + "--steps", + str(steps), + "-W", + str(generation_width), + "-H", + str(generation_height), + "--upscale-model", + str(upscale_model.expanduser()), + "--upscale-repeats", + str(upscale_repeats), + "--upscale-tile-size", + str(upscale_tile_size), + "-p", + prompt_text, + "-o", + str(output_path), + ] + + if flash_attention: + cmd.append("--fa") + + logger.info( + "Diffuse at %dx%d", + generation_width, + generation_height, + ) + try: - perma = store / "generated_wallpaper.jpg" - _ = perma.open("wb").write(Path(saved[idx]).read_bytes()) - logger.info(f"Setting wallpaper to {perma}") - _ = subprocess.run(command + [perma], check=True) - except Exception as e: - logger.warning(f"Could not set wallpaper: {e}") + subprocess.run(cmd, check=True) + except FileNotFoundError as exc: + raise RuntimeError("Could not find the sd-cli executable") from exc + except subprocess.CalledProcessError as exc: + raise RuntimeError(f"sd-cli exited with status {exc.returncode}") from exc + + if not output_path.is_file(): + raise RuntimeError( + f"sd-cli completed successfully but did not create {output_path}" + ) + + prompt_path.write_text(prompt_text, encoding="utf-8") + return output_path + + +def step_wallpaper(command: list[str], saved: Path, store: Path) -> None: + try: + store.mkdir(parents=True, exist_ok=True) + permanent_path = store / "generated_wallpaper.png" + shutil.copyfile(saved, permanent_path) + + logger.info("Setting wallpaper to %s", permanent_path) + subprocess.run( + [*command, str(permanent_path)], + check=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + logger.warning("Could not set wallpaper: %s", exc) def parse_args() -> argparse.Namespace: @@ -382,6 +364,7 @@ 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", @@ -405,10 +388,12 @@ def parse_args() -> argparse.Namespace: ) wall.add_argument( "--taste-sample-frac", - type=float, + type=positive_fraction, default=0.8, - help="Fraction of taste keywords randomly sampled per run; lower " - "values give more varied, less cluttered prompts.", + help=( + "Fraction of taste keywords randomly sampled per run; lower " + "values give more varied, less cluttered prompts." + ), ) wall.add_argument( "--biomes", @@ -423,8 +408,10 @@ def parse_args() -> argparse.Namespace: "tropical highlands: lush jungle terraces, waterfalls down cliff faces, climbing vines, humid haze, broad-leafed plants", "megastructure interior: vast enclosed halls, endless repeating bays, shafts of light from distant openings, no vegetation", ], - help="Biome descriptions; one is chosen at random per run and " - "injected as the dominant environment.", + help=( + "Biome descriptions; one is chosen at random per run and " + "injected as the dominant environment." + ), ) wall.add_argument( "--perspectives", @@ -440,22 +427,22 @@ def parse_args() -> argparse.Namespace: "view from flat open ground, structures rising ahead", "framed view through a narrow gap or corridor", ], - help="Camera perspectives; one is chosen at random per run and " - "injected as a hard directive.", + 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." + "--ephemeral-dir", + type=Path, + default=Path("/tmp/wallpaper"), + help="Output directory.", ) wall.add_argument( "--single-file-persist-path", - default="~/media/pictures/", - help="Single target for persistance.", - ) - wall.add_argument( - "--target-width", type=int, default=3840, help="Final wallpaper width." - ) - wall.add_argument( - "--target-height", type=int, default=2160, help="Final wallpaper height." + type=Path, + default=Path("~/media/pictures"), + help="Directory for the persistent wallpaper file.", ) wall.add_argument( "--set-command", @@ -463,43 +450,42 @@ def parse_args() -> argparse.Namespace: default=["awww", "img"], help="Command to run when setting wallpaper.", ) + llm = parser.add_argument_group("llm") llm.add_argument( - "--prompt-only", action="store_true", help="Terminate after running the LLM." - ) - llm.add_argument( - "--llm-repo-id", - default="unsloth/gemma-4-E4B-it-GGUF", - help="HuggingFace repo for LLM to use to prompt image model.", + "--prompt-only", + action="store_true", + help="Print the generated prompt and terminate before diffusion.", ) llm.add_argument( - "--llm-filename", - default="gemma-4-E4B-it-Q4_K_M.gguf", - help="File name of LLM in the repo to use.", + "--llm-model-path", + type=Path, + default=Path("~/store/weights/Qwen3-8B-Q4_K_M.gguf"), + help="Model path for generating image prompts.", ) llm.add_argument( "--llm-max-tokens", - type=int, + type=positive_int, default=384, - help="Max tokens of prompt to generate.", + help="Maximum number of prompt tokens to generate.", ) llm.add_argument( "--llm-temperature", - type=float, + type=non_negative_float, default=0.4, - help="Sampling temperature; llama-cpp-python's default is 0.2", + help="Sampling temperature.", ) llm.add_argument( "--llm-top-p", - type=float, + type=probability, default=0.95, help="Nucleus sampling threshold.", ) llm.add_argument( "--llm-min-p", - type=float, + type=probability, default=0.05, - help="Minimum token probability relative to the top token", + help="Minimum token probability relative to the top token.", ) llm.add_argument( "--llm-system-prompt", @@ -519,69 +505,102 @@ def parse_args() -> argparse.Namespace: "Do not include the time, date, or location in your prompt. " "Output ONLY the image prompt — no explanation or preamble." ), - help="System prompt for LLM.", - ) - 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.", - ) - flux2.add_argument( - "--flux2-klein-filename", - default="flux-2-klein-4b-Q8_0.gguf", - help="File name of quantised mode in the repo.", - ) - flux2.add_argument( - "--flux2-klein-steps", - type=int, - default=5, - help="Steps to use when generating image.", - ) - flux2.add_argument( - "--flux2-klein-gen-width", type=int, default=1280, help="Generated image width." - ) - flux2.add_argument( - "--flux2-klein-gen-height", - type=int, + help="System prompt for the LLM.", + ) + + sd_cli = parser.add_argument_group("sd-cli") + sd_cli.add_argument( + "--sd-diffusion-model", + type=Path, + default=Path("~/store/weights/flux-2-klein-9b-Q4_0.gguf"), + help="Diffusion model passed to --diffusion-model.", + ) + sd_cli.add_argument( + "--sd-vae", + type=Path, + default=Path("~/store/weights/vae.safetensors"), + help="VAE model passed to --vae.", + ) + sd_cli.add_argument( + "--sd-llm", + type=Path, + default=Path("~/store/weights/Qwen3-8B-Q4_K_M.gguf"), + help="Text encoder or LLM passed to --llm.", + ) + sd_cli.add_argument( + "--sd-cfg-scale", + type=non_negative_float, + default=1.0, + help="Guidance scale passed to --cfg-scale.", + ) + sd_cli.add_argument( + "--sd-steps", + type=positive_int, + default=4, + help="Sampling steps passed to --steps.", + ) + sd_cli.add_argument( + "--sd-flash-attention", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable or disable the sd-cli --fa option.", + ) + sd_cli.add_argument( + "--sd-width", + type=positive_int, + default=1280, + help="Base diffusion generation width passed to -W.", + ) + sd_cli.add_argument( + "--sd-height", + type=positive_int, default=720, - help="Generated image height.", - ) - sr = parser.add_argument_group("super-resolution") - sr.add_argument( - "--sr-repo-id", - default="Phips/4xNomos2_hq_mosr", - help="HuggingFace repo hosting a spandrel-supported checkpoint.", - ) - sr.add_argument( - "--sr-filename", - default="4xNomos2_hq_mosr.safetensors", - help="Checkpoint filename in the repo, must be spandrel supported.", - ) - sr.add_argument( - "--sr-tile", - type=int, - default=1024, - help="Tile size for tiled upscaling; bounds peak memory on constrained machines. 0 = full-frame (best quality).", + help="Base diffusion generation height passed to -H.", + ) + sd_cli.add_argument( + "--sd-upscale-model", + type=Path, + default=Path("~/store/weights/4xNomosWebPhoto_esrgan.safetensors"), + help="ESRGAN model passed to --upscale-model.", + ) + sd_cli.add_argument( + "--sd-upscale-repeats", + type=non_negative_int, + default=1, + help="Number of ESRGAN passes passed to --upscale-repeats.", + ) + sd_cli.add_argument( + "--sd-upscale-tile-size", + type=positive_int, + default=256, + help="Tile size passed to --upscale-tile-size.", + ) + sd_cli.add_argument( + "--sd-cli-explicit-instructions", + default="No railings, no glass, no people, no objects.", + help="Instructions appended to every diffusion prompt.", ) - wx = parser.add_argument_group("weather") - wx.add_argument( - "--latitude", type=float, default=51.75, help="Latitude for weather." + + weather = parser.add_argument_group("weather") + weather.add_argument( + "--latitude", + type=float, + default=51.75, + help="Latitude for weather.", ) - wx.add_argument( - "--longitude", type=float, default=-1.25, help="Longitude for weather." + weather.add_argument( + "--longitude", + type=float, + default=-1.25, + help="Longitude for weather.", ) - wx.add_argument( + weather.add_argument( "--elevation", type=int, default=61, - help="Altitude (m) for sunrise/sunset times.", + help="Altitude in metres for sunrise and sunset times.", ) + return parser.parse_args() @@ -593,83 +612,57 @@ def run() -> None: stream=sys.stderr, ) args = parse_args() + now = dt.datetime.now().astimezone() - # crazy - now = dt.datetime.now(dt.datetime.now().astimezone().tzinfo) - logger.info(f"Capability {CAP}, dtype {COMPUTE_DTYPE}") - - # 1. context - ctx = step_context(now, args.latitude, args.longitude, args.elevation) - logger.info(f"Context: {ctx}") + context = step_context( + now, + args.latitude, + args.longitude, + args.elevation, + ) + logger.info("Context: %s", context) - # 2. prompt rewrite via local LLM prompt_text = step_llm( - args.llm_repo_id, - args.llm_filename, - args.llm_system_prompt, - args.llm_max_tokens, - list(args.taste_prompt), - args.taste_sample_frac, - list(args.biomes), - list(args.perspectives), - args.llm_temperature, - args.llm_top_p, - args.llm_min_p, - ctx, - args.flux2_explicit_instructions, - ) - logger.info(f"Prompt: {prompt_text}") + system_prompt=args.llm_system_prompt, + model_path=args.llm_model_path, + max_tokens=args.llm_max_tokens, + taste_prompts=list(args.taste_prompt), + taste_sample_frac=args.taste_sample_frac, + biomes=list(args.biomes), + perspectives=list(args.perspectives), + temperature=args.llm_temperature, + top_p=args.llm_top_p, + min_p=args.llm_min_p, + context=context, + explicit_instructions=args.sd_cli_explicit_instructions, + ) + logger.info("Prompt: %s", prompt_text) if args.prompt_only: + print(prompt_text) return - # 3. diffusion at base resolution - logger.info( - "Diffuse: %dx%d @ %d steps", - args.flux2_klein_gen_width, - args.flux2_klein_gen_height, - args.flux2_klein_steps, - ) - raw_images = step_diffuse( - args.flux2_klein_repo_id, - args.flux2_klein_filename, - args.flux2_klein_gen_width, - args.flux2_klein_gen_height, - args.flux2_klein_steps, - prompt_text, + saved = step_diffuse( + diffusion_model=args.sd_diffusion_model, + vae=args.sd_vae, + llm=args.sd_llm, + cfg_scale=args.sd_cfg_scale, + steps=args.sd_steps, + flash_attention=args.sd_flash_attention, + generation_width=args.sd_width, + generation_height=args.sd_height, + upscale_model=args.sd_upscale_model, + upscale_repeats=args.sd_upscale_repeats, + upscale_tile_size=args.sd_upscale_tile_size, + output_dir=args.ephemeral_dir.expanduser(), + now=now, + prompt_text=prompt_text, ) - # 4. super-resolution then scale - logger.info( - "Upscale then rescale: %dx%d -> upscale -> %dx%d via %s with tiling %d", - args.flux2_klein_gen_width, - args.flux2_klein_gen_height, - args.target_width, - args.target_height, - args.sr_filename, - args.sr_tile, - ) - upscaled = step_upscale( - raw_images, - args.target_width, - args.target_height, - args.sr_repo_id, - args.sr_filename, - args.sr_tile, - ) - - # 5. save to configured output directory - saved = step_save( - upscaled, raw_images, 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, - Path(args.single_file_persist_path).expanduser(), + command=args.set_command, + saved=saved, + store=args.single_file_persist_path.expanduser(), ) |