diff options
Diffstat (limited to 'wallpaper.py')
| -rwxr-xr-x | wallpaper.py | 128 |
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, |