tinygrab/examples/vgg7.py

import sys
import random
import json
import numpy
from pathlib import Path
from PIL import Image
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import SGD
from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict
from examples.vgg7_helpers.waifu2x import image_load, image_save, Vgg7

# amount of context erased by model
CONTEXT = 7


def get_sample_count(samples_dir):
    try:
        samples_dir_count_file = open(samples_dir + "/sample_count.txt", "r")
        v = samples_dir_count_file.readline()
        samples_dir_count_file.close()
        return int(v)
    except:
        return 0


def set_sample_count(samples_dir, sc):
    with open(samples_dir + "/sample_count.txt", "w") as file:
        file.write(str(sc) + "\n")


if len(sys.argv) < 2:
    print("python3 -m examples.vgg7 import MODELJSON MODEL")
    print(
        " imports a waifu2x JSON vgg_7 model, i.e. waifu2x/models/vgg_7/art/scale2.0x_model.json"
    )
    print(" into a safetensors file")
    print(" weight tensors are ordered in tinygrad/ncnn form, as so: (outC,inC,H,W)")
    print(" *this format is used by most other commands in this program*")
    print("python3 -m examples.vgg7 import_kinne MODEL_KINNE MODEL_SAFETENSORS")
    print(
        " imports a model in 'KINNE' format (raw floats: used by older versions of this example) into safetensors"
    )
    print("python3 -m examples.vgg7 execute MODEL IMG_IN IMG_OUT")
    print(" given an already-nearest-neighbour-scaled image, runs vgg7 on it")
    print(" output image has 7 pixels removed on all edges")
    print(" do not run on large images, will have *hilarious* RAM use")
    print("python3 -m examples.vgg7 execute_full MODEL IMG_IN IMG_OUT")
    print(" does the 'whole thing' (padding, tiling)")
    print(" safe for large images, etc.")
    print("python3 -m examples.vgg7 new MODEL")
    print(" creates a new model (experimental)")
    print("python3 -m examples.vgg7 train MODEL SAMPLES_DIR ROUNDS ROUNDS_SAVE")
    print(" trains a model (experimental)")
    print(" (how experimental? well, every time I tried it, it flooded w/ NaNs)")
    print(" note: ROUNDS < 0 means 'forever'. ROUNDS_SAVE <= 0 is not a good idea.")
    print(" expects roughly execute's input as SAMPLES_DIR/IDXa.png")
    print(" expects roughly execute's output as SAMPLES_DIR/IDXb.png")
    print(" (i.e. my_samples/0a.png is the first pre-nearest-scaled image,")
    print("       my_samples/0b.png is the first original image)")
    print(" in addition, SAMPLES_DIR/samples_count.txt indicates sample count")
    print(" won't pad or tile, so keep image sizes sane")
    print("python3 -m examples.vgg7 samplify IMG_A IMG_B SAMPLES_DIR SIZE")
    print(
        " creates overlapping micropatches (SIZExSIZE w/ 7-pixel border) for training"
    )
    print(" maintains/creates samples_count.txt automatically")
    print(" unlike training, IMG_A must be exactly half the size of IMG_B")
    sys.exit(1)

cmd = sys.argv[1]
vgg7 = Vgg7()


def nansbane(p):
    if numpy.isnan(numpy.min(p.numpy())):
        raise Exception(
            "A NaN in the model has been detected. This model will not be interacted with to prevent further damage."
        )


def load_and_save(path, save):
    if save:
        for v in vgg7.get_parameters():
            nansbane(v)
        st = get_state_dict(vgg7)
        safe_save(st, path)
    else:
        st = safe_load(path)
        load_state_dict(vgg7, st)
        for v in vgg7.get_parameters():
            nansbane(v)


if cmd == "import":
    src = sys.argv[2]
    model = sys.argv[3]

    vgg7.load_waifu2x_json(json.load(open(src, "rb")))

    load_and_save(model, True)
elif cmd == "import_kinne":
    # tinygrad wasn't doing safetensors when this example was written
    # it's possible someone might have a model around using the resulting interim format
    src = sys.argv[2]
    model = sys.argv[3]

    index = 0
    for t in vgg7.get_parameters():
        fn = src + "/snoop_bin_" + str(index) + ".bin"
        t.assign(Tensor(numpy.fromfile(fn, "<f4")).reshape(shape=t.shape))
        index += 1

    load_and_save(model, True)
elif cmd == "execute":
    model = sys.argv[2]
    in_file = sys.argv[3]
    out_file = sys.argv[4]

    load_and_save(model, False)

    image_save(out_file, vgg7.forward(Tensor(image_load(in_file))).numpy())
elif cmd == "execute_full":
    model = sys.argv[2]
    in_file = sys.argv[3]
    out_file = sys.argv[4]

    load_and_save(model, False)

    image_save(out_file, vgg7.forward_tiled(image_load(in_file), 156))
elif cmd == "new":
    model = sys.argv[2]

    load_and_save(model, True)
elif cmd == "train":
    model = sys.argv[2]
    samples_base = sys.argv[3]
    samples_count = get_sample_count(samples_base)
    rounds = int(sys.argv[4])
    rounds_per_save = int(sys.argv[5])

    load_and_save(model, False)

    # Initialize sample probabilities.
    # This is used to try and get the network to focus on "interesting" samples,
    #  which works nicely with the microsample system.
    sample_probs = None
    sample_probs_path = model + "_sample_probs.bin"
    try:
        # try to read...
        sample_probs = numpy.fromfile(sample_probs_path, "<f8")
        if sample_probs.shape[0] != samples_count:
            print("sample probs size != sample count - initializing")
            sample_probs = None
    except:
        # it's fine
        print("sample probs could not be loaded - initializing")

    if sample_probs is None:
        # This stupidly high amount is used to force an initial pass over all samples
        sample_probs = numpy.ones(samples_count) * 1000

    print("Training...")
    # Adam has a tendency to destroy the state of the network when restarted
    # Plus it's slower
    optim = SGD(vgg7.get_parameters())

    rnum = 0
    while True:
        # The way the -1 option works is that rnum is never -1.
        if rnum == rounds:
            break

        sample_idx = 0
        try:
            sample_idx = numpy.random.choice(
                samples_count, p=sample_probs / sample_probs.sum()
            )
        except:
            print("exception occurred (PROBABLY value-probabilities-dont-sum-to-1)")
            sample_idx = random.randint(0, samples_count - 1)

        x_img = image_load(samples_base + "/" + str(sample_idx) + "a.png")
        y_img = image_load(samples_base + "/" + str(sample_idx) + "b.png")

        sample_x = Tensor(x_img, requires_grad=False)
        sample_y = Tensor(y_img, requires_grad=False)

        # magic code roughly from readme example
        # An explanation, in case anyone else has to go down this path:
        # This runs the actual network normally
        out = vgg7.forward(sample_x)
        # Subtraction determines error here (as this is an image, not classification).
        # *Abs is the important bit* - at least for me, anyway.
        # The training process seeks to minimize this 'loss' value.
        # Minimization of loss *tends towards negative infinity*, so without the abs,
        #  or without an implicit abs (the mul in the README),
        #  loss will always go haywire in one direction or another.
        # Mean determines how errors are treated.
        # Do not use Sum. I tried that. It worked while I was using 1x1 patches...
        # Then it went exponential.
        # Also, Mean goes *after* abs. I realize this should have been obvious to me.
        loss = sample_y.sub(out).abs().mean()
        # This is the bit where tinygrad works backward from the loss
        optim.zero_grad()
        loss.backward()
        # And this updates the parameters
        optim.step()

        # warning: used by sample probability adjuster
        loss_indicator = loss.max().numpy()
        print("Round " + str(rnum) + " : " + str(loss_indicator))

        if (rnum % rounds_per_save) == 0:
            print("Saving")
            load_and_save(model, True)
            sample_probs.astype("<f8", "C").tofile(sample_probs_path)

        # Update round state
        # Number
        rnum = rnum + 1
        # Probability management
        # there must always be a probability, no matter how slim, even if loss goes to 0
        sample_probs[sample_idx] = max(loss_indicator, 1.0e-10)

    # if we were told to save every round, we already saved
    if rounds_per_save != 1:
        print("Done with all rounds, saving")
        load_and_save(model, True)
        sample_probs.astype("<f8", "C").tofile(sample_probs_path)

elif cmd == "samplify":
    a_img = sys.argv[2]
    b_img = sys.argv[3]
    samples_base = sys.argv[4]
    sample_size = int(sys.argv[5])
    samples_count = get_sample_count(samples_base)

    # This bit is interesting because it actually does some work.
    # Not much, but some work.
    a_img = image_load(a_img)
    b_img = image_load(b_img)

    # as with the main library body,
    # Y X order is used here

    # assertion before pre-upscaling is performed
    assert a_img.shape[2] == (b_img.shape[2] // 2)
    assert a_img.shape[3] == (b_img.shape[3] // 2)

    # pre-upscaling - this matches the sizes (and coordinates)
    a_img = a_img.repeat(2, 2).repeat(2, 3)

    samples_added = 0

    # actual patch extraction
    for posy in range(
        CONTEXT, b_img.shape[2] - (CONTEXT + sample_size - 1), sample_size
    ):
        for posx in range(
            CONTEXT, b_img.shape[3] - (CONTEXT + sample_size - 1), sample_size
        ):
            # this is a viable patch location, add it
            # note the ranges here:
            #  + there are always CONTEXT pixels *before* the point
            #  + with no subtraction at the end, there'd already be a pixel *at* the point,
            #     as ranges are exclusive
            #  + additionally, there are sample_size - 1 additional sample pixels
            #  + additionally, there are CONTEXT additional pixels
            #  + therefore there are CONTEXT + sample_size pixels *at & after* the point
            patch_x = a_img[
                :,
                :,
                posy - CONTEXT : posy + CONTEXT + sample_size,
                posx - CONTEXT : posx + CONTEXT + sample_size,
            ]
            patch_y = b_img[:, :, posy : posy + sample_size, posx : posx + sample_size]

            image_save(f"{samples_base}/{str(samples_count)}a.png", patch_x)
            image_save(f"{samples_base}/{str(samples_count)}b.png", patch_y)
            samples_count += 1
            samples_added += 1

    print(f"Added {str(samples_added)} samples")
    set_sample_count(samples_base, samples_count)

else:
    print("unknown command")