diff options
Diffstat (limited to 'wallpaper.py')
| -rwxr-xr-x | wallpaper.py | 554 |
1 files changed, 554 insertions, 0 deletions
diff --git a/wallpaper.py b/wallpaper.py new file mode 100755 index 0000000..b2048b2 --- /dev/null +++ b/wallpaper.py @@ -0,0 +1,554 @@ +#!/usr/bin/env -S uv run --script +# +# /// 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 subprocess +import sys +import gc + +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 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 + + +def step_context( + now: dt.datetime, + latitude: float | None, + longitude: float | None, + elevation: int = 0, +) -> str: + + today = now.strftime("Current date and time: %A %B %d, %Y at %H:%M %Z") + + 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 = ( + 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", + } + try: + url = ( + f"https://api.open-meteo.com/v1/forecast?" + f"latitude={latitude}&longitude={longitude}" + f"¤t=weather_code,temperature_2m&timezone=auto" + ) + 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: + logger.warning("Weather fetch failed (%s), continuing anyway.", exc) + + return f"{today}. Daylight: {sun_times}. Weather: {detail}" if detail else today + + +def step_llm( + repo_id: str, + filename: str, + system_prompt: str, + max_tokens: int, + taste_prompts: list[str], + 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) + + prompt = ( + f"Taste reference:\n{taste_prompt_str}\n\n" + f"Context: {ctx}\n\n" + f"Reply with one text-to-image prompt only." + ) + logger.info("Running llm...") + output = llm.create_chat_completion( + messages=[ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": prompt, + }, + ], + max_tokens=max_tokens, + ) + + del llm + gc.collect() + return output["choices"][0]["message"]["content"].strip() + + +def step_diffuse( + repo_id: str, + filename: str, + gen_width: int, + gen_height: int, + 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, + ) + 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, + ) + 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, + ) + + 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("cpu").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) + + 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( + images: list[Any], output_dir: Path, now: dt.datetime, prompt_text: str +) -> list[Path]: + output_dir.mkdir(parents=True, exist_ok=True) + saved = [] + tag = now.strftime("%Y%m%d-%H%M%S") + for i, img in enumerate(images): + fname = f"wallpaper-{tag}" + (f"-{i}" if len(images) > 1 else "") + ".jpg" + 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 + + +def step_wallpaper(command: list[str], saved: list[Path]) -> None: + idx = random.randrange(len(saved)) + try: + logger.info(f"Setting wallpaper to {saved[idx]}") + _ = subprocess.run(command + [saved[idx]], check=True) + except Exception as e: + logger.warning(f"Could not set wallpaper: {e}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + 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", + nargs="+", + default=[ + "natural scenes", + "serene", + "trees", + "plants", + "landscape", + "hypermodern concrete architecture", + "moss", + "rocks", + "abstract concrete sculptures", + "water channel", + "distant futuristic megastructure", + ], + help="Taste keywords for the LLM art-director prompt.", + ) + wall.add_argument( + "--output-dir", default="/tmp/wallpaper/", help="Output directory." + ) + 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." + ) + wall.add_argument( + "--set-command", + nargs="+", + default=["awww", "img"], + help="Command to run when setting wallpaper.", + ) + + llm = parser.add_argument_group("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.", + ) + 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.add_argument( + "--llm-max-tokens", + type=int, + default=200, + help="Max tokens of prompt to generate.", + ) + llm.add_argument( + "--llm-system-prompt", + default=( + "You are a wallpaper art director. Given the user's taste and today's context, " + "write one detailed text-to-image prompt for a diffusion model. Not all of the " + "user's taste should be present, it's a suggestion and you should choose only some " + "or expand on it as required. The time of day should feature in your output clearly. " + "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. " + "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-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=4, + 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, + default=720, + help="Generated image height.", + ) + + sr = parser.add_argument_group("super-resolution") + sr.add_argument( + "--sr-repo-id", + default="Comfy-Org/Real-ESRGAN_repackaged", + help="HuggingFace repo hosting the ESRGAN-family checkpoint.", + ) + sr.add_argument( + "--sr-filename", + default="RealESRGAN_x4plus.safetensors", + help="Checkpoint filename in the repo (any spandrel-supported " + ".pth/.safetensors works).", + ) + sr.add_argument( + "--sr-tile", + type=int, + default=256, + 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." + ) + wx.add_argument( + "--longitude", type=float, default=-1.25, help="Longitude for weather." + ) + wx.add_argument( + "--elevation", + type=int, + default=61, + help="Altitude (m) for sunrise/sunset times.", + ) + + return parser.parse_args() + + +def run() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(message)s", + 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 + ctx = step_context(now, args.latitude, args.longitude, args.elevation) + logger.info(f"Context: {ctx}") + + # 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), + ctx, + ) + logger.info(f"Prompt: {prompt_text}") + + # 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, + ) + + # 4. super-resolution then scale + logger.info( + "Upscale then rescale: %dx%d -> %dx%d -> %dx%d via %s with tiling %d", + args.flux2_klein_gen_width, + args.flux2_klein_gen_height, + args.flux2_klein_gen_width * 4, + args.flux2_klein_gen_height * 4, + 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, Path(args.output_dir).expanduser(), now, prompt_text) + + for p in saved: + logger.info("Saved: %s", p.resolve()) + + # 6. set wallpaper + step_wallpaper(args.set_command, saved) + + +if __name__ == "__main__": + run() |