File size: 5,813 Bytes
ed46c5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Diffusion-based synthetic photographs of runic inscriptions (latent diffusion + ControlNet-Canny).
Follows the recipe of the recognition thesis: the transliteration is mapped to runes, rendered
black on white on a 1024x1024 canvas with automatic font size and line wrapping, converted to a
Canny edge map (thresholds 50/150, one-pixel dilation) and used as the ControlNet condition with
a fixed positive/negative prompt pair, 28 steps and guidance 7.0. Adapted for SDXL: a 1216x832 canvas instead of
1024x1024 and conditioning scale 0.8 instead of 0.65, with extra negative terms, because SDXL fills empty
stone with invented signs at the thesis settings.
SDXL-base with the SDXL Canny ControlNet replaces SD3-Medium, whose weights are gated.
Resumable per shard; labels are appended after every image.
"""
import argparse
import json
import os
import random
import sys
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from synth import FONT_DIR, RUNE, lexicon, sample_phrase # noqa: E402
POSITIVE = ("deeply carved runes on a dark grey granite runestone, weathered eroded stone surface, "
"plain stone around the inscription, even diffuse daylight, archaeological macro photograph, sharp focus")
NEGATIVE = ("polished, smooth, printed font, fantasy, illustration, painting, cartoon, glowing, "
"text overlay, watermark, signature, blurry, extra runes, additional inscriptions, symbols, ornaments, "
"patterns, tiles, bricks, frame")
def render_condition(translit, rng, width=1216, height=832):
runes = " ".join("".join(RUNE.get(c, "") for c in w) for w in translit.split())
runes = runes.replace(" ", " ᛫ ") # runic single punctuation as word divider
font_path = os.path.join(FONT_DIR, "NotoSansRunic-Regular.ttf")
size = width
canvas = Image.new("L", (width, height), 255)
draw = ImageDraw.Draw(canvas)
fsize = 200
while fsize > 24:
font = ImageFont.truetype(font_path, fsize)
words, lines, cur = runes.split(" "), [], ""
for w in words:
trial = (cur + " " + w).strip()
if draw.textlength(trial, font=font) <= size * 0.86:
cur = trial
else:
lines.append(cur)
cur = w
lines.append(cur)
line_h = int(fsize * 1.35)
if len(lines) * line_h <= height * 0.8 and all(draw.textlength(l, font=font) <= size * 0.86 for l in lines):
break
fsize -= 8
y0 = (height - len(lines) * line_h) // 2
for i, l in enumerate(lines):
w = draw.textlength(l, font=font)
x = (size - w) / 2 + rng.randint(-20, 20)
draw.text((x, y0 + i * line_h), l, fill=0, font=font)
edges = cv2.Canny(np.array(canvas), 50, 150)
edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1)
return Image.fromarray(edges).convert("RGB")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--corpus", required=True)
ap.add_argument("--out_dir", required=True)
ap.add_argument("--shard", type=int, default=0)
ap.add_argument("--n", type=int, default=1000)
ap.add_argument("--scale", type=float, default=0.8)
args = ap.parse_args()
import torch
from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetPipeline
hub = "/home/jovyan/.cache/huggingface/hub" # job pods have a different $HOME
snap = lambda repo: sorted(os.listdir(f"{hub}/models--{repo.replace('/', '--')}/snapshots"))[0] # noqa: E731
base = f"{hub}/models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/{snap('stabilityai/stable-diffusion-xl-base-1.0')}"
cn = f"{hub}/models--diffusers--controlnet-canny-sdxl-1.0/snapshots/{snap('diffusers/controlnet-canny-sdxl-1.0')}"
vae = f"{hub}/models--madebyollin--sdxl-vae-fp16-fix/snapshots/{snap('madebyollin/sdxl-vae-fp16-fix')}"
controlnet = ControlNetModel.from_pretrained(cn, torch_dtype=torch.float16, variant="fp16")
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
base, controlnet=controlnet, vae=AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16),
torch_dtype=torch.float16, variant="fp16").to("cuda")
pipe.set_progress_bar_config(disable=True)
rng = random.Random(1000 + args.shard)
lex = lexicon(args.corpus)
words = sorted(lex)
weights = [lex[w] ** 0.7 for w in words]
os.makedirs(os.path.join(args.out_dir, "img"), exist_ok=True)
label_path = os.path.join(args.out_dir, f"labels_{args.shard}.jsonl")
done = set()
if os.path.exists(label_path):
done = {json.loads(l)["id"] for l in open(label_path)}
phrases = [sample_phrase(rng, words, weights) for _ in range(args.n)]
with open(label_path, "a") as f:
for i, t in enumerate(phrases):
uid = f"dif{args.shard}_{i:05d}" # v2 parameters (1216x832 canvas, scale 0.8)
if uid in done:
continue
cond = render_condition(t, rng)
g = torch.Generator("cuda").manual_seed(args.shard * 100000 + i)
img = pipe(POSITIVE, negative_prompt=NEGATIVE, image=cond, width=cond.width, height=cond.height, num_inference_steps=28, guidance_scale=7.0,
controlnet_conditioning_scale=args.scale, generator=g).images[0]
img.thumbnail((768, 768))
name = f"img/{uid}.jpg"
img.save(os.path.join(args.out_dir, name), quality=90)
f.write(json.dumps({"id": uid, "image": name, "target": t}, ensure_ascii=False) + "\n")
f.flush()
if i % 50 == 0:
print("generated", i, flush=True)
print("DONE", args.shard, flush=True)
if __name__ == "__main__":
main()
|