Sana Coloring

Colors black-and-white manga pages. You can let the model choose the colors, or steer the scene like an artist with five simple choices. It is built on Sana and was trained on about 60 GB of manga pages across genres.

Controls

Each control can be left empty (<null>): the model then chooses by itself. When every control is empty you get the model's own unprompted colors.

• mood: calm, tense, horror, sad, action, comedic, romantic, nostalgic, mysterious, cheerful, melancholic, dramatic
• time_of_day: morning, dawn, day, evening, night
• lighting: soft diffuse, golden warm, neutral, strong contrast, flat gray-blue, golden orange, warm orange-yellow, cool blue, cool silver, pale harsh, magenta-cyan
• saturation: muted, moderate, vivid
• palette: cold blue-gray with a touch of dull red, warm amber and brown, soft pastel pinks and blues, soft pastel pink and light blue, deep teal and orange, green and earthy brown, green and brown, purple and midnight blue, purple and dark navy, bright primary colors, bright red, blue and yellow, sepia and cream, sepia brown and cream, crimson and black, crimson red and black, fresh spring greens and sky blue, light green and light blue

Results

How to use

Put data.py, model.py, ema_bf16.pt and meta.json from this repo in one folder. The weights of the Sana base model are downloaded the first time Colorizer is built.

import json, torch
from PIL import Image
from torchvision.transforms import v2
from torchvision.transforms.functional import to_pil_image
from data import FIELDS, TAGS
from model import Colorizer, load_vae, sample

meta = json.load(open("meta.json"))
model = Colorizer(meta["kept"]).to("cuda").eval()
model.load_state_dict(torch.load("ema_bf16.pt"))
model.forward = torch.autocast("cuda", torch.bfloat16)(model.forward)
vae = load_vae(device="cuda")

page = Image.open("page.png").convert("L").convert("RGB")
scale = 512 / min(page.size)
h, w = (max(32, round(d * scale / 32) * 32) for d in page.size[::-1])
bw = v2.functional.to_dtype(v2.functional.resize(v2.functional.pil_to_tensor(page), [h, w], antialias=True), torch.float32, scale=True)
bw = ((bw - 0.5) / 0.5)[None].cuda()

tags = {f: "<null>" for f in FIELDS}
tags.update(saturation="vivid", palette="bright primary colors")        # leave everything out for "model chooses"
ids = torch.tensor([[TAGS[f].str2int(v) for f, v in tags.items()]]).cuda()

out = sample(model, vae, bw, ids, 20, 2.0, 3.0, 3.0 * 512 / 512)
to_pil_image(out[0].clamp(0, 1).cpu()).save("colored.png")

Training details

This checkpoint is the EMA weights of the teacher phase at step 30000, trained at 512 px. The full list of training settings is in params.json.

Limitations

• Very dense hatching and screentone can turn into blotchy patches.
• Clothing and skin colors are sometimes confused.
• The colors follow the model's training data and tend to be soft unless you pick vivid tags.