KohakuBlueleaf

"""Minimal, self-contained reproduction of the FloorSet-Lite MIB inconsistency.

MIB (Multiple-Instantiation Block) requires every member of a group to share IDENTICAL (w,h)
dimensions -- which forces an identical area. This script proves, straight from IntelLabs' published
data (no project code on the train path), that the TRAIN set assigns members of one MIB group
DIFFERENT target areas, so the MIB constraint is unsatisfiable by construction. It then runs the exact
same check on the VAL set, which is clean -- demonstrating the check is correct and the bug is in the
train data, not the check.

Raw layout tensor layout (FloorSet-Lite): ``c[0][i]`` is ``[n,6]`` with column 0 = target area and

KohakuBlueleaf

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
import torchvision.models as models
from torch.utils.data import Dataset
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

KohakuBlueleaf

"""
An Minimal Implementation of IMM (Inductive Moment Matching)
"""

import math
import torch
import torch.nn.functional as F


def compute_mmd_loss_fully_vectorized(

KohakuBlueleaf

## Note
## if use vllm in same gpu, remember to set a low gpu_memory_utilization to avoid OOM
## For larger model please consider to use multi-GPU or CPU offloading
## AnySchedule: https://github.com/KohakuBlueleaf/AnySchedule
## LyCORIS: https://github.com/KohakuBlueleaf/LyCORIS
## Following code can perform reasonable training on Llama-3.2-1B-Instruct model with GSM8K dataset
## With noticable improvement on each reward function
from itertools import chain
import re

KohakuBlueleaf

Sd-WebUI standalone script

just throw the script into the root folder of sd-webui. and then modify the content to what you want to run. And then run it.

If you want to use some extension, just use my implementation as example (to modify the args for extensions)

KohakuBlueleaf

import torch
import torch.nn as nn
import torch.nn.functional as F


class TimestepEmbedding(nn.Module):
    def __init__(self, n_time=1024, n_emb=320, n_out=1280) -> None:
        super().__init__()
        self.emb = nn.Embedding(n_time, n_emb)

KohakuBlueleaf

import asyncio
import cv2
import json
import traceback
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import numpy as np
from httpx import AsyncClient, HTTPError, Timeout
from tqdm import tqdm

KohakuBlueleaf

import torch
import torch.nn as nn
import torch.nn.functional as F

from einops import rearrange


def parallel_retention(
    q, k, v,            # bsz, heads, seq_len, dim
    decay_mask = None   # heads, seq_len, seq_len

KohakuBlueleaf

import os, sys

from time import time_ns
from json import load
from io import BytesIO
import math

from PIL import Image

import asyncio

KohakuBlueleaf

import torch
import numpy as np
import timeit

#could be "cpu" "cuda" "opencl" "mps"...
DEVICE = "mps"
ROUNDS = 500
NUM = (2048, 2048)

# Input for benchmarking