src/compliment_forest/backends/image.py

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)