from __future__ import annotations import base64 import hashlib import io import random from pathlib import Path from typing import Protocol from PIL import Image, ImageDraw, ImageFilter, ImageOps STYLE_PREFIX = ( "cmprst_forest, soft watercolor storybook illustration, loose wet-on-wet washes, " "visible cold-press paper grain, no hard outlines, soft feathered edges, " ) STYLE_SUFFIX = ( ", warm dappled light, kind expression, lots of negative space, muted warm palette" ) NEGATIVE_PROMPT = ( "hard outlines, photorealistic, 3d render, neon, oversaturated, busy background, " "text, watermark, deformed" ) class ImageBackend(Protocol): def generate(self, prompt: str, seed: int) -> str: ... def compose_flux_prompt(creature_prompt: str) -> str: return f"{STYLE_PREFIX}{creature_prompt.strip()}{STYLE_SUFFIX}" def image_to_data_uri(image: Image.Image) -> str: buffer = io.BytesIO() image.save(buffer, format="PNG", optimize=True) encoded = base64.b64encode(buffer.getvalue()).decode("ascii") return f"data:image/png;base64,{encoded}" class DemoImageBackend: """Creates stable local artwork for tests and UI development.""" def __init__( self, width: int = 768, height: int = 576, asset_dir: str | Path | None = None, ) -> None: self.width = width self.height = height self.asset_dir = ( Path(asset_dir) if asset_dir is not None else Path(__file__).resolve().parents[3] / "frontend" / "assets" ) def _curated_asset(self, prompt: str, seed: int) -> Image.Image | None: prompt_lower = prompt.casefold() names = { "owl": "owl.png", "snail": "snail.png", "deer": "deer.png", "wren": "wren.png", "fox": "fox-path.png", } filename = next( (asset for keyword, asset in names.items() if keyword in prompt_lower), "fox-path.png", ) path = self.asset_dir / filename if not path.exists(): return None with Image.open(path) as source: image = ImageOps.fit( source.convert("RGB"), (self.width, self.height), method=Image.Resampling.LANCZOS, ) corner = image.getpixel((0, 0)) variation = seed % 3 image.putpixel((0, 0), tuple(min(255, channel + variation) for channel in corner)) return image def generate(self, prompt: str, seed: int) -> str: curated = self._curated_asset(prompt, seed) if curated is not None: return image_to_data_uri(curated) prompt_seed = int.from_bytes( hashlib.sha256(prompt.encode("utf-8")).digest()[:4], byteorder="big", ) rng = random.Random(seed ^ prompt_seed) image = Image.new("RGB", (self.width, self.height), "#f5f0df") wash = Image.new("RGBA", image.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(wash, "RGBA") palette = [ (124, 145, 112, 28), (77, 105, 78, 24), (208, 151, 139, 22), (225, 184, 101, 22), (132, 172, 181, 18), ] for _ in range(100): color = rng.choice(palette) x = rng.randint(-100, self.width) y = rng.randint(-80, self.height) radius_x = rng.randint(35, 160) radius_y = rng.randint(25, 110) draw.ellipse( (x - radius_x, y - radius_y, x + radius_x, y + radius_y), fill=color, ) wash = wash.filter(ImageFilter.GaussianBlur(radius=22)) image = Image.alpha_composite(image.convert("RGBA"), wash) creature = Image.new("RGBA", image.size, (0, 0, 0, 0)) creature_draw = ImageDraw.Draw(creature, "RGBA") center_x = self.width // 2 + rng.randint(-30, 30) ground_y = int(self.height * 0.72) body_color = rng.choice( [(169, 105, 69, 205), (128, 105, 78, 205), (154, 135, 96, 205)] ) shadow = (76, 91, 67, 35) creature_draw.ellipse( (center_x - 145, ground_y - 10, center_x + 145, ground_y + 38), fill=shadow, ) creature_draw.ellipse( (center_x - 96, ground_y - 170, center_x + 90, ground_y + 5), fill=body_color, ) creature_draw.ellipse( (center_x - 72, ground_y - 245, center_x + 74, ground_y - 100), fill=body_color, ) creature_draw.polygon( [ (center_x - 62, ground_y - 210), (center_x - 88, ground_y - 290), (center_x - 20, ground_y - 235), ], fill=body_color, ) creature_draw.polygon( [ (center_x + 58, ground_y - 210), (center_x + 88, ground_y - 290), (center_x + 20, ground_y - 235), ], fill=body_color, ) creature_draw.ellipse( (center_x - 38, ground_y - 195, center_x - 15, ground_y - 169), fill=(51, 53, 44, 220), ) creature_draw.ellipse( (center_x + 18, ground_y - 195, center_x + 41, ground_y - 169), fill=(51, 53, 44, 220), ) creature_draw.ellipse( (center_x - 6, ground_y - 163, center_x + 9, ground_y - 150), fill=(72, 57, 47, 190), ) creature = creature.filter(ImageFilter.GaussianBlur(radius=1.2)) image = Image.alpha_composite(image, creature) grain = Image.new("RGBA", image.size, (0, 0, 0, 0)) grain_pixels = grain.load() for y in range(self.height): for x in range(self.width): value = rng.randint(0, 10) grain_pixels[x, y] = (75, 66, 50, value) image = Image.alpha_composite(image, grain) return image_to_data_uri(image.convert("RGB")) class FluxImageBackend: """Lazy FLUX.1-dev pipeline that loads a local or Hub LoRA.""" def __init__( self, model_id: str = "black-forest-labs/FLUX.1-dev", lora_id: str = "build-small-hackathon/compliment-forest-flux-lora", *, local_files_only: bool = False, width: int = 768, height: int = 768, steps: int = 28, ) -> None: self.model_id = model_id self.lora_id = lora_id self.local_files_only = local_files_only self.width = width self.height = height self.steps = steps self._pipeline = None def _load(self): if self._pipeline is not None: return self._pipeline import torch from diffusers import FluxPipeline pipeline = FluxPipeline.from_pretrained( self.model_id, torch_dtype=torch.bfloat16, local_files_only=self.local_files_only, ) pipeline.load_lora_weights( self.lora_id, local_files_only=self.local_files_only, ) pipeline.enable_model_cpu_offload() self._pipeline = pipeline return pipeline def generate(self, prompt: str, seed: int) -> str: import torch pipeline = self._load() generator = torch.Generator(device="cpu").manual_seed(seed) image = pipeline( prompt=compose_flux_prompt(prompt), negative_prompt=NEGATIVE_PROMPT, width=self.width, height=self.height, num_inference_steps=self.steps, guidance_scale=3.5, generator=generator, ).images[0] return image_to_data_uri(image)