Add application file

app.py

@@ -0,0 +1,514 @@
+import gradio as gr
+import pandas as pd
+import numpy as np
+import torch as th
+from house_diffusion import dist_util
+from house_diffusion.script_util import (
+    create_model_and_diffusion,
+)
+from PIL import Image
+
+import io
+import drawSvg as drawsvg
+import cairosvg
+from tqdm import tqdm
+
+import webcolors
+import tempfile
+from pathlib import Path
+import shutil
+import os
+
+
+ROOM_CLASS = {
+    'Living Room': 1, 'Kitchen': 2, 'Bedroom': 3, 'Bathroom': 4,
+    'Balcony': 5, 'Entrance': 6, 'Dining Room': 7, 'Study Room': 8,
+    'Storage': 10, 'Front Door': 11, 'Unknown': 13, 'Interior Door': 12
+}
+ROOM_CATEGORIES = {
+    'Living Room': 1, 'Kitchen': 2, 'Bedroom': 3, 'Bathroom': 4,
+    'Balcony': 5, 'Entrance': 6, 'Dining Room': 7, 'Study Room': 8,
+    'Storage': 10, 'Front Door': 11, 'Other': 13
+}
+
+
+def save_samples(
+        sample, ext, model_kwargs,
+        tmp_count, num_room_types,
+        # save_gif=False,
+        save_gif=True,
+        door_indices=[11, 12, 13], ID_COLOR=None,
+        is_syn=False, draw_graph=False, save_svg=False, metrics=False):
+    prefix = 'syn_' if is_syn else ''
+    graph_errors = []
+
+    print(sample.shape)
+
+    if not save_gif:
+        sample = sample[-1:]
+    for i in tqdm(range(sample.shape[1])):
+        resolution = 256
+        images = []
+        images2 = []
+        images3 = []
+        for k in range(sample.shape[0]):
+            draw_color = drawsvg.Drawing(resolution, resolution, displayInline=False)
+            draw_color.append(drawsvg.Rectangle(0, 0, resolution, resolution, fill='white'))
+            polys = []
+            types = []
+            for j, point in (enumerate(sample[k][i])):
+                if model_kwargs[f'{prefix}src_key_padding_mask'][i][j] == 1:
+                    continue
+                point = point.cpu().data.numpy()
+                if j == 0:
+                    poly = []
+                if j > 0 and (model_kwargs[f'{prefix}room_indices'][i, j] != model_kwargs[f'{prefix}room_indices'][
+                    i, j - 1]).any():
+                    polys.append(poly)
+                    types.append(c)
+                    poly = []
+                pred_center = False
+                if pred_center:
+                    point = point / 2 + 1
+                    point = point * resolution // 2
+                else:
+                    point = point / 2 + 0.5
+                    point = point * resolution
+                poly.append((point[0], point[1]))
+                c = np.argmax(model_kwargs[f'{prefix}room_types'][i][j - 1].cpu().numpy())
+            polys.append(poly)
+            types.append(c)
+            for poly, c in zip(polys, types):
+                if c in door_indices or c == 0:
+                    continue
+                room_type = c
+                c = webcolors.hex_to_rgb(ID_COLOR[c])
+                draw_color.append(
+                    drawsvg.Lines(*np.array(poly).flatten().tolist(), close=True, fill=ID_COLOR[room_type],
+                                  fill_opacity=1.0, stroke='black', stroke_width=1))
+
+            for poly, c in zip(polys, types):
+                if c not in door_indices:
+                    continue
+                room_type = c
+                c = webcolors.hex_to_rgb(ID_COLOR[c])
+
+                # TODO --------------------------------------------------------------------------------------
+                # https://github.com/sakmalh/house_diffusion
+                line_lengths = [np.linalg.norm(np.array(poly[i]) - np.array(poly[(i + 1) % len(poly)])) for i in
+                                range(len(poly))]
+
+                if metrics:
+                    text_size = 5
+                    for z, length in enumerate(line_lengths):
+                        # Calculate the mid-point of the line segment
+                        midpoint = ((poly[z][0] + poly[(z + 1) % len(poly)][0]) / 2,
+                                    (poly[z][1] + poly[(z + 1) % len(poly)][1]) / 2)
+
+                        # Calculate x and y differences
+                        x_diff = poly[z][0] - poly[(z + 1) % len(poly)][0]
+                        y_diff = poly[z][1] - poly[(z + 1) % len(poly)][1]
+
+                        # Determine text position adjustments based on differences
+                        if int(y_diff) != 0:
+                            if y_diff > 0:
+                                text_x = midpoint[0] + text_size
+                                text_y = midpoint[1]
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x, text_y + text_size,  # Start point at the text label
+                                            poly[z][0] + text_size, poly[z][1],  # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x, text_y - text_size,  # Start point at the text label
+                                            poly[(z + 1) % len(poly)][0] + text_size, poly[(z + 1) % len(poly)][1],
+                                    # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+                            else:
+                                text_x = midpoint[0] - text_size
+                                text_y = midpoint[1]
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x, text_y - text_size,  # Start point at the text label
+                                            poly[z][0] - text_size, poly[z][1],  # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x, text_y + text_size,  # Start point at the text label
+                                            poly[(z + 1) % len(poly)][0] - text_size, poly[(z + 1) % len(poly)][1],
+                                    # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+                        else:
+                            if x_diff > 0:
+                                text_x = midpoint[0]
+                                text_y = midpoint[1] - text_size
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x + text_size, text_y,  # Start point at the text label
+                                    poly[z][0], poly[z][1] - text_size,  # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x - text_size, text_y,  # Start point at the text label
+                                    poly[(z + 1) % len(poly)][0], poly[(z + 1) % len(poly)][1] - text_size,
+                                    # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+                            else:
+                                text_x = midpoint[0]
+                                text_y = midpoint[1] + text_size
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x - text_size, text_y,  # Start point at the text label
+                                    poly[z][0], poly[z][1] + text_size,  # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+
+                                draw_color.append(drawsvg.Line(
+                                    text_x + text_size, text_y,  # Start point at the text label
+                                    poly[(z + 1) % len(poly)][0], poly[(z + 1) % len(poly)][1] + text_size,
+                                    # End point at the polygon endpoint
+                                    stroke='black',
+                                    stroke_width=1
+                                ))
+
+                        # Add the text label to the SVG
+                        draw_color.append(
+                            drawsvg.Text(
+                                f'{int(abs(length))}',  # Format the length to two decimal places
+                                text_size,
+                                text_x, text_y,
+                                fill='black',
+                                text_anchor='middle',
+                                alignment_baseline='middle'
+                            )
+                        )
+
+                draw_color.append(
+                    drawsvg.Lines(*np.array(poly).flatten().tolist(), close=True, fill=ID_COLOR[room_type],
+                                  fill_opacity=1.0, stroke='black', stroke_width=1))
+
+            if k == sample.shape[0] - 1 or True:
+                if save_svg:
+                    # draw_color.saveSvg(f'outputs/{ext}/{tmp_count + i}c_{k}_{ext}.svg')
+                    return draw_color
+                else:
+                    Image.open(io.BytesIO(cairosvg.svg2png(draw_color.asSvg()))).save(
+                        f'outputs/{ext}/{tmp_count + i}c_{ext}.png')
+
+        # if save_gif:
+        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}.gif', images, fps=10, loop=1)
+        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}_v2.gif', images2, fps=10, loop=1)
+        #     imageio.mimwrite(f'outputs/gif/{tmp_count + i}_v3.gif', images3, fps=10, loop=1)
+    return graph_errors
+
+
+def function_test(org_graphs, corners, room_type):
+    get_one_hot = lambda x, z: np.eye(z)[x]
+    max_num_points = 100
+
+    house = []
+    corner_bounds = []
+    num_points = 0
+
+    for i, room in enumerate(room_type):
+        # Adding conditions
+        num_room_corners = corners[i]
+        rtype = np.repeat(np.array([get_one_hot(room, 25)]), num_room_corners, 0)
+        room_index = np.repeat(np.array([get_one_hot(len(house) + 1, 32)]), num_room_corners, 0)
+        corner_index = np.array([get_one_hot(x, 32) for x in range(num_room_corners)])
+        # Src_key_padding_mask
+        padding_mask = np.repeat(1, num_room_corners)
+        padding_mask = np.expand_dims(padding_mask, 1)
+        # Generating corner bounds for attention masks
+        connections = np.array([[i, (i + 1) % num_room_corners] for i in range(num_room_corners)])
+        connections += num_points
+        corner_bounds.append([num_points, num_points + num_room_corners])
+        num_points += num_room_corners
+        room = np.concatenate((np.zeros([num_room_corners, 2]), rtype, corner_index, room_index,
+                               padding_mask, connections), 1)
+        house.append(room)
+
+    house_layouts = np.concatenate(house, 0)
+    padding = np.zeros((max_num_points - len(house_layouts), 94))
+    gen_mask = np.ones((max_num_points, max_num_points))
+    gen_mask[:len(house_layouts), :len(house_layouts)] = 0
+    house_layouts = np.concatenate((house_layouts, padding), 0)
+
+    door_mask = np.ones((max_num_points, max_num_points))
+    self_mask = np.ones((max_num_points, max_num_points))
+    for i, room in enumerate(room_type):
+        if room == 1:
+            living_room_index = i
+            break
+    for i in range(len(corner_bounds)):
+        is_connected = False
+        for j in range(len(corner_bounds)):
+            if i == j:
+                self_mask[corner_bounds[i][0]:corner_bounds[i][1], corner_bounds[j][0]:corner_bounds[j][1]] = 0
+            elif any(np.equal([i, 1, j], org_graphs).all(1)) or any(np.equal([j, 1, i], org_graphs).all(1)):
+                door_mask[corner_bounds[i][0]:corner_bounds[i][1], corner_bounds[j][0]:corner_bounds[j][1]] = 0
+                is_connected = True
+        if not is_connected:
+            door_mask[corner_bounds[i][0]:corner_bounds[i][1],
+            corner_bounds[living_room_index][0]:corner_bounds[living_room_index][1]] = 0
+
+    syn_houses = house_layouts
+    syn_door_masks = door_mask
+    syn_self_masks = self_mask
+    syn_gen_masks = gen_mask
+
+    syn_graph = np.concatenate((org_graphs, np.zeros([200 - len(org_graphs), 3])), 0)
+
+    cond = {
+        'syn_door_mask': syn_door_masks,
+        'syn_self_mask': syn_self_masks,
+        'syn_gen_mask': syn_gen_masks,
+        'syn_room_types': syn_houses[:, 2:2 + 25],
+        'syn_corner_indices': syn_houses[:, 2 + 25:2 + 57],
+        'syn_room_indices': syn_houses[:, 2 + 57:2 + 89],
+        'syn_src_key_padding_mask': 1 - syn_houses[:, 2 + 89],
+        'syn_connections': syn_houses[:, 2 + 90:2 + 92],
+        'syn_graph': syn_graph,
+    }
+
+    return cond
+
+
+def create_layout(graphs, corners, room_type, metrics=False, use_ddim=True, ddim_steps=100, num_samples=4):
+    model_path = "ckpt/model250000.pt"
+    steps = f"ddim{ddim_steps}"
+    args = {
+        "input_channels": 18,
+        "condition_channels": 89,
+        "num_channels": 512,
+        "out_channels": 2,
+        "dataset": "rplan",
+        "use_checkpoint": False,
+        "use_unet": False,
+        "learn_sigma": False,
+        "diffusion_steps": 1000,
+        "noise_schedule": "cosine",
+        "timestep_respacing": steps,
+        "use_kl": False,
+        "predict_xstart": False,
+        "rescale_timesteps": False,
+        "rescale_learned_sigmas": False,
+        "analog_bit": False,
+        "target_set": -1,
+        "set_name": "",
+    }
+
+    dist_util.setup_dist()
+    model, diffusion = create_model_and_diffusion(
+        args['input_channels'],
+        args['condition_channels'],
+        args['num_channels'],
+        args['out_channels'],
+        args['dataset'],
+        args['use_checkpoint'],
+        args['use_unet'],
+        args['learn_sigma'],
+        args['diffusion_steps'],
+        args['noise_schedule'],
+        args['timestep_respacing'],
+        args['use_kl'],
+        args['predict_xstart'],
+        args['rescale_timesteps'],
+        args['rescale_learned_sigmas'],
+        args['analog_bit'],
+        args['target_set'],
+        args['set_name'],
+    )
+    model.load_state_dict(
+        dist_util.load_state_dict(model_path, map_location="cpu")
+    )
+    model.to(dist_util.dev())
+    model.eval()
+    ID_COLOR = {1: '#EE4D4D', 2: '#C67C7B', 3: '#FFD274', 4: '#BEBEBE', 5: '#BFE3E8',
+                6: '#7BA779', 7: '#E87A90', 8: '#FF8C69', 10: '#1F849B', 11: '#727171',
+                13: '#785A67', 12: '#D3A2C7'}
+    num_room_types = 14
+    sample_fn = (diffusion.p_sample_loop if not use_ddim else diffusion.ddim_sample_loop)
+    print(graphs, corners, room_type)
+    model_kwargs = function_test(graphs, corners, room_type)
+    for key in model_kwargs:
+        model_kwargs[key] = th.from_numpy(np.array([model_kwargs[key]])).cuda()
+
+    png_paths = []
+    svg_paths = []
+    for count in range(num_samples):
+        sample = sample_fn(
+            model,
+            th.Size([1, 2, 100]),
+            clip_denoised=True,
+            model_kwargs=model_kwargs,
+        )
+
+        sample = sample.permute([0, 1, 3, 2])
+
+        pred = save_samples(sample, 'pred', model_kwargs, count, num_room_types, ID_COLOR=ID_COLOR,
+                            is_syn=True, draw_graph=False, save_svg=True, save_gif=False, metrics=metrics)
+
+        temp_svg_file = tempfile.NamedTemporaryFile(delete=False, suffix=".svg")
+        pred.saveSvg(temp_svg_file.name)
+        png_file_name = temp_svg_file.name.split(".")[0].split("/")[-1]
+        png_file_path = f'./generated_svgs/{png_file_name}.png'
+        # print(temp_svg_file.name)
+        # print(png_file_name)
+        # print(png_file_path)
+
+        Image.open(io.BytesIO(cairosvg.svg2png(pred.asSvg()))).save(png_file_path)
+
+        output_dir = Path("./generated_svgs")
+        output_dir.mkdir(parents=True, exist_ok=True)
+        file_name = temp_svg_file.name.split("/")[-1]
+        persistent_path = Path(f"{output_dir}/{file_name}")
+        shutil.move(temp_svg_file.name, persistent_path)
+        os.chmod(persistent_path, 0o644)
+
+        svg_paths.append(str(persistent_path))
+        png_paths.append(png_file_path)
+        # print(str(persistent_path))
+
+    return png_paths, svg_paths
+
+
+rooms_data = []
+edges_data = []
+
+
+def generate_layout(metrics: bool, ddim_steps: int, num_samples: int):
+    room_list = []
+    room_corners = []
+    living_room = 0
+    front_door = False
+    entrance = -1
+
+    print(rooms_data)
+    print(edges_data)
+
+    for i, room in enumerate(rooms_data):
+        room_list.append(ROOM_CLASS[room['room_type']])
+        if room['num_corners'] != 0:
+            room_corners.append(int(room['num_corners']))
+        else:
+            room_corners.append(4)
+
+        if room['room_type'] == "Living Room":
+            living_room = i
+
+        elif room['room_type'] == "Entrance":
+            entrance = i
+
+        elif room['room_type'] == "Front Door":
+            front_door = True
+
+    edges = []
+    for edge in edges_data:
+        source_id = int(edge['room1_id'].split()[0])
+        target_id = int(edge['room2_id'].split()[0])
+        edges.append([source_id, 1, target_id])
+
+        index = len(room_list)
+        room_list.append(12)
+        room_corners.append(4)
+        edges.append([source_id, 1, index])
+        edges.append([target_id, 1, index])
+
+    if not front_door:
+        room_list.append(11)
+        room_corners.append(4)
+        if entrance == -1:
+            edges.append([len(room_list) - 1, 1, living_room])
+        else:
+            edges.append([len(room_list) - 1, 1, entrance])
+
+    if np.sum(room_corners) > 99:
+        return {"Error": "Number of Corners exceeded"}
+
+    print(room_list, room_corners, edges)
+    png_paths, svg_paths = create_layout(edges, room_corners, room_list, metrics=metrics, ddim_steps=ddim_steps,
+                             num_samples=num_samples)
+
+    png_color_guide = './color_guide.png'
+
+    return png_paths, svg_paths, png_color_guide
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## House Layout Generator")
+
+    with gr.Row():
+        room_type = gr.Dropdown(label="Room Type", choices=list(ROOM_CATEGORIES.keys()), value="Living Room")
+        num_corners = gr.Number(label="Number of Corners", value=4)
+        add_room_button = gr.Button("Add Room")
+
+    with gr.Row():
+        room1_id = gr.Dropdown(label="Room 1", choices=[], value=None)
+        room2_id = gr.Dropdown(label="Room 2", choices=[], value=None)
+        add_edge_button = gr.Button("Add Edge")
+
+    rooms_table = gr.DataFrame(label="Rooms Table")
+    edges_table = gr.DataFrame(label="Edges Table")
+
+    metrics_toggle = gr.Checkbox(label="Include metrics", value=True)
+    ddim_input = gr.Number(label="DDIM steps", value=100)
+    num_sample = gr.Number(label="Number of samples", value=4)
+
+    png_gallery = gr.Gallery(label="Layout PNG Outputs", columns=4)
+    svg_files = gr.File(label="Layout SVG Outputs (higher quality)")
+    png_color_guide = gr.Image(label="Color Guide")
+
+    def add_room(room_type, num_corners):
+        room_id = len(rooms_data)
+        rooms_data.append({
+            "room_id": room_id,
+            "room_type": room_type,
+            "num_corners": num_corners
+        })
+        return update_rooms_and_edges()
+
+
+    def add_edge(room1_id, room2_id):
+        edge_id = len(edges_data)
+        edges_data.append({
+            "edge_id": edge_id,
+            "room1_id": room1_id,
+            "room2_id": room2_id
+        })
+        return update_rooms_and_edges()
+
+
+    def update_rooms_and_edges():
+        rooms_df = pd.DataFrame(rooms_data, columns=["room_id", "room_type", "num_corners"])
+        edges_df = pd.DataFrame(edges_data, columns=["edge_id", "room1_id", "room2_id"])
+        room_options = [f"{room['room_id']} {room['room_type']}" for room in rooms_data]
+        return rooms_df, edges_df, gr.update(choices=room_options, value=None), gr.update(choices=room_options,
+                                                                                          value=None)
+
+
+    generate_button = gr.Button("Generate Layout")
+    generate_button.click(generate_layout, inputs=[metrics_toggle, ddim_input, num_sample], outputs=[png_gallery, svg_files, png_color_guide])
+
+    add_room_button.click(add_room, inputs=[room_type, num_corners],
+                          outputs=[rooms_table, edges_table, room1_id, room2_id])
+    add_edge_button.click(add_edge, inputs=[room1_id, room2_id], outputs=[rooms_table, edges_table, room1_id, room2_id])
+
+
+demo.launch()
+# global demo
+# demo.launch(share=True)

ckpt/model250000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b22c916ec8d5fce087ce9bfe277817a93d60ca2eed9f8bbcf4b8eefec43797a
+size 106240205

house_diffusion/__init__.py

@@ -0,0 +1,3 @@
+"""
+Codebase for "HouseDiffusion" based on the implementation from "Improved Denoising Diffusion Probabilistic Models".
+"""

house_diffusion/dist_util.py

@@ -0,0 +1,94 @@
+"""
+Helpers for distributed training.
+"""
+
+import io
+import os
+import socket
+
+import blobfile as bf
+from mpi4py import MPI
+import torch as th
+import torch.distributed as dist
+
+# Change this to reflect your cluster layout.
+# The GPU for a given rank is (rank % GPUS_PER_NODE).
+GPUS_PER_NODE = 4
+
+SETUP_RETRY_COUNT = 3
+
+
+def setup_dist():
+    """
+    Setup a distributed process group.
+    """
+    if dist.is_initialized():
+        return
+    ## temporary removed to manually set the CUDA_VISIBLE_DEVICES
+    #os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}"
+
+    comm = MPI.COMM_WORLD
+    backend = "gloo" if not th.cuda.is_available() else "nccl"
+
+    if backend == "gloo":
+        hostname = "localhost"
+    else:
+        hostname = socket.gethostbyname(socket.getfqdn())
+    os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0)
+    os.environ["RANK"] = str(comm.rank)
+    os.environ["WORLD_SIZE"] = str(comm.size)
+
+    port = comm.bcast(_find_free_port(), root=0)
+    os.environ["MASTER_PORT"] = str(port)
+    dist.init_process_group(backend=backend, init_method="env://")
+
+
+def dev():
+    """
+    Get the device to use for torch.distributed.
+    """
+    if th.cuda.is_available():
+        return th.device(f"cuda")
+    return th.device("cpu")
+
+
+def load_state_dict(path, **kwargs):
+    """
+    Load a PyTorch file without redundant fetches across MPI ranks.
+    """
+    chunk_size = 2 ** 30  # MPI has a relatively small size limit
+    if MPI.COMM_WORLD.Get_rank() == 0:
+        with bf.BlobFile(path, "rb") as f:
+            data = f.read()
+        num_chunks = len(data) // chunk_size
+        if len(data) % chunk_size:
+            num_chunks += 1
+        MPI.COMM_WORLD.bcast(num_chunks)
+        for i in range(0, len(data), chunk_size):
+            MPI.COMM_WORLD.bcast(data[i : i + chunk_size])
+    else:
+        num_chunks = MPI.COMM_WORLD.bcast(None)
+        data = bytes()
+        for _ in range(num_chunks):
+            data += MPI.COMM_WORLD.bcast(None)
+
+    return th.load(io.BytesIO(data), **kwargs)
+
+
+def sync_params(params):
+    """
+    Synchronize a sequence of Tensors across ranks from rank 0.
+    """
+    for p in params:
+        with th.no_grad():
+            dist.broadcast(p, 0)
+
+
+def _find_free_port():
+    try:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+    finally:
+        s.close()

house_diffusion/fp16_util.py

@@ -0,0 +1,236 @@
+"""
+Helpers to train with 16-bit precision.
+"""
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from . import logger
+
+INITIAL_LOG_LOSS_SCALE = 20.0
+
+
+def convert_module_to_f16(l):
+    """
+    Convert primitive modules to float16.
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.half()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.half()
+
+
+def convert_module_to_f32(l):
+    """
+    Convert primitive modules to float32, undoing convert_module_to_f16().
+    """
+    if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        l.weight.data = l.weight.data.float()
+        if l.bias is not None:
+            l.bias.data = l.bias.data.float()
+
+
+def make_master_params(param_groups_and_shapes):
+    """
+    Copy model parameters into a (differently-shaped) list of full-precision
+    parameters.
+    """
+    master_params = []
+    for param_group, shape in param_groups_and_shapes:
+        master_param = nn.Parameter(
+            _flatten_dense_tensors(
+                [param.detach().float() for (_, param) in param_group]
+            ).view(shape)
+        )
+        master_param.requires_grad = True
+        master_params.append(master_param)
+    return master_params
+
+
+def model_grads_to_master_grads(param_groups_and_shapes, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    for master_param, (param_group, shape) in zip(
+        master_params, param_groups_and_shapes
+    ):
+        master_param.grad = _flatten_dense_tensors(
+            [param_grad_or_zeros(param) for (_, param) in param_group]
+        ).view(shape)
+
+
+def master_params_to_model_params(param_groups_and_shapes, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    # Without copying to a list, if a generator is passed, this will
+    # silently not copy any parameters.
+    for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes):
+        for (_, param), unflat_master_param in zip(
+            param_group, unflatten_master_params(param_group, master_param.view(-1))
+        ):
+            param.detach().copy_(unflat_master_param)
+
+
+def unflatten_master_params(param_group, master_param):
+    return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group])
+
+
+def get_param_groups_and_shapes(named_model_params):
+    named_model_params = list(named_model_params)
+    scalar_vector_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim <= 1],
+        (-1),
+    )
+    matrix_named_params = (
+        [(n, p) for (n, p) in named_model_params if p.ndim > 1],
+        (1, -1),
+    )
+    return [scalar_vector_named_params, matrix_named_params]
+
+
+def master_params_to_state_dict(
+    model, param_groups_and_shapes, master_params, use_fp16
+):
+    if use_fp16:
+        state_dict = model.state_dict()
+        for master_param, (param_group, _) in zip(
+            master_params, param_groups_and_shapes
+        ):
+            for (name, _), unflat_master_param in zip(
+                param_group, unflatten_master_params(param_group, master_param.view(-1))
+            ):
+                assert name in state_dict
+                state_dict[name] = unflat_master_param
+    else:
+        state_dict = model.state_dict()
+        for i, (name, _value) in enumerate(model.named_parameters()):
+            assert name in state_dict
+            state_dict[name] = master_params[i]
+    return state_dict
+
+
+def state_dict_to_master_params(model, state_dict, use_fp16):
+    if use_fp16:
+        named_model_params = [
+            (name, state_dict[name]) for name, _ in model.named_parameters()
+        ]
+        param_groups_and_shapes = get_param_groups_and_shapes(named_model_params)
+        master_params = make_master_params(param_groups_and_shapes)
+    else:
+        master_params = [state_dict[name] for name, _ in model.named_parameters()]
+    return master_params
+
+
+def zero_master_grads(master_params):
+    for param in master_params:
+        param.grad = None
+
+
+def zero_grad(model_params):
+    for param in model_params:
+        # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group
+        if param.grad is not None:
+            param.grad.detach_()
+            param.grad.zero_()
+
+
+def param_grad_or_zeros(param):
+    if param.grad is not None:
+        return param.grad.data.detach()
+    else:
+        return th.zeros_like(param)
+
+
+class MixedPrecisionTrainer:
+    def __init__(
+        self,
+        *,
+        model,
+        use_fp16=False,
+        fp16_scale_growth=1e-3,
+        initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE,
+    ):
+        self.model = model
+        self.use_fp16 = use_fp16
+        self.fp16_scale_growth = fp16_scale_growth
+
+        self.model_params = list(self.model.parameters())
+        self.master_params = self.model_params
+        self.param_groups_and_shapes = None
+        self.lg_loss_scale = initial_lg_loss_scale
+
+        if self.use_fp16:
+            self.param_groups_and_shapes = get_param_groups_and_shapes(
+                self.model.named_parameters()
+            )
+            self.master_params = make_master_params(self.param_groups_and_shapes)
+            self.model.convert_to_fp16()
+
+    def zero_grad(self):
+        zero_grad(self.model_params)
+
+    def backward(self, loss: th.Tensor):
+        if self.use_fp16:
+            loss_scale = 2 ** self.lg_loss_scale
+            (loss * loss_scale).backward()
+        else:
+            loss.backward()
+
+    def optimize(self, opt: th.optim.Optimizer):
+        if self.use_fp16:
+            return self._optimize_fp16(opt)
+        else:
+            return self._optimize_normal(opt)
+
+    def _optimize_fp16(self, opt: th.optim.Optimizer):
+        logger.logkv_mean("lg_loss_scale", self.lg_loss_scale)
+        model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params)
+        grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale)
+        if check_overflow(grad_norm):
+            self.lg_loss_scale -= 1
+            logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}")
+            zero_master_grads(self.master_params)
+            return False
+
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+
+        self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale))
+        opt.step()
+        zero_master_grads(self.master_params)
+        master_params_to_model_params(self.param_groups_and_shapes, self.master_params)
+        self.lg_loss_scale += self.fp16_scale_growth
+        return True
+
+    def _optimize_normal(self, opt: th.optim.Optimizer):
+        grad_norm, param_norm = self._compute_norms()
+        logger.logkv_mean("grad_norm", grad_norm)
+        logger.logkv_mean("param_norm", param_norm)
+        opt.step()
+        return True
+
+    def _compute_norms(self, grad_scale=1.0):
+        grad_norm = 0.0
+        param_norm = 0.0
+        for p in self.master_params:
+            with th.no_grad():
+                param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2
+                if p.grad is not None:
+                    grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2
+        return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm)
+
+    def master_params_to_state_dict(self, master_params):
+        return master_params_to_state_dict(
+            self.model, self.param_groups_and_shapes, master_params, self.use_fp16
+        )
+
+    def state_dict_to_master_params(self, state_dict):
+        return state_dict_to_master_params(self.model, state_dict, self.use_fp16)
+
+
+def check_overflow(value):
+    return (value == float("inf")) or (value == -float("inf")) or (value != value)

house_diffusion/gaussian_diffusion.py

@@ -0,0 +1,1013 @@
+"""
+This code started out as a PyTorch port of Ho et al's diffusion models:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
+
+Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
+"""
+
+import enum
+import math
+
+import numpy as np
+import torch as th
+
+from .nn import mean_flat
+from .losses import normal_kl, discretized_gaussian_log_likelihood
+from tqdm.auto import tqdm
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == "cosine":
+        print("COSINE")
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            # lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+            lambda t: math.cos((t) / 1.000 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here, and then adapted over time to further experimentation.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    :param model_mean_type: a ModelMeanType determining what the model outputs.
+    :param model_var_type: a ModelVarType determining how variance is output.
+    :param loss_type: a LossType determining the loss function to use.
+    :param rescale_timesteps: if True, pass floating point timesteps into the
+                              model so that they are always scaled like in the
+                              original paper (0 to 1000).
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type,
+        rescale_timesteps=False,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.rescale_timesteps = rescale_timesteps
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # log calculation clipped because the posterior variance is 0 at the
+        # beginning of the diffusion chain.
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev)
+            * np.sqrt(alphas)
+            / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        )
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(
+            self.log_one_minus_alphas_cumprod, t, x_start.shape
+        )
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+            * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None, analog_bit=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        xtalpha = _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape).permute([0,2,1])
+        epsalpha = _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape).permute([0,2,1])
+        # print("TTTTTTTTTTTTTTTTTTTT", t)
+        model_output_dec, model_output_bin = model(x, self._scale_timesteps(t), xtalpha=xtalpha, epsalpha=epsalpha, is_syn=True, **model_kwargs)
+        model_output = model_output_dec
+
+        if analog_bit:
+            predict_descrete = 0
+        else:
+            predict_descrete = 32
+
+        if t[0] < predict_descrete:
+            def bin2dec(b, bits):
+                mask = 2 ** th.arange(bits - 1, -1, -1).to(b.device, b.dtype)
+                return th.sum(mask * b, -1)
+            model_output_bin[model_output_bin>0] = 1
+            model_output_bin[model_output_bin<=0] = 0
+            model_output_bin = bin2dec(model_output_bin.round().int().permute([0,2,1]).reshape(model_output_bin.shape[0],
+                model_output_bin.shape[2], 2, 8), 8).permute([0,2,1])
+
+            model_output_bin = ((model_output_bin/256) - 0.5) * 2
+            model_output = model_output_bin
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == ModelVarType.LEARNED:
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(
+                    self.posterior_log_variance_clipped, t, x.shape
+                )
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if t[0] >= predict_descrete:
+            if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+                )
+                model_mean = model_output
+            elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
+                if self.model_mean_type == ModelMeanType.START_X:
+                    pred_xstart = process_xstart(model_output)
+                else:
+                    pred_xstart = process_xstart(
+                        self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                    )
+                model_mean, _, _ = self.q_posterior_mean_variance(
+                    x_start=pred_xstart, x_t=x, t=t
+                )
+            else:
+                raise NotImplementedError(self.model_mean_type)
+        else:
+            pred_xstart = process_xstart(model_output)
+            model_mean, _, _ = self.q_posterior_mean_variance(
+                x_start=pred_xstart, x_t=x, t=t
+                )
+
+        assert (
+            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        )
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _scale_timesteps(self, t):
+        if self.rescale_timesteps:
+            return t.float() * (1000.0 / self.num_timesteps)
+        return t
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
+        new_mean = (
+            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        )
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
+            x, self._scale_timesteps(t), **model_kwargs
+        )
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(
+            x_start=out["pred_xstart"], x_t=x, t=t
+        )
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        analog_bit=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+            analog_bit=analog_bit,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(
+                cond_fn, out, x, t, model_kwargs=model_kwargs
+            )
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        analog_bit=None,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        myfinal = []
+        final = None
+        for i, sample in tqdm(enumerate(self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            analog_bit=analog_bit,
+        ))):
+            if i>970:
+                myfinal.append(sample['sample'])
+            final = sample
+        return th.stack(myfinal)
+        # return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        analog_bit=None,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    analog_bit=analog_bit,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_next)
+            + th.sqrt(1 - alpha_bar_next) * eps
+        )
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        # old code
+        # myfinal = []
+        # for i, sample in tqdm(enumerate(self.ddim_sample_loop_progressive(
+        #     model,
+        #     shape,
+        #     noise=noise,
+        #     clip_denoised=clip_denoised,
+        #     denoised_fn=denoised_fn,
+        #     cond_fn=cond_fn,
+        #     model_kwargs=model_kwargs,
+        #     device=device,
+        #     progress=progress,
+        #     eta=eta,
+        # ))):
+        #     # if i>998:
+        #     if i>498:
+        #     # if i>98:
+        #         myfinal.append(sample['sample'])
+        # return th.stack(myfinal)
+
+        # final = None
+        # for sample in self.ddim_sample_loop_progressive(
+        #         model,
+        #         shape,
+        #         noise=noise,
+        #         clip_denoised=clip_denoised,
+        #         denoised_fn=denoised_fn,
+        #         cond_fn=cond_fn,
+        #         model_kwargs=model_kwargs,
+        #         device=device,
+        #         progress=progress,
+        #         eta=eta,
+        # ):
+        #     final = sample
+        # return final["sample"]
+        
+        samples_list = list(tqdm(self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+        )))
+
+        last_sample = samples_list[-1]['sample']
+
+        return th.stack([last_sample])
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+        self, model, x_start, x_t, t, padding_mask, clip_denoised=True, model_kwargs=None,
+    ):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl, padding_mask) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll, padding_mask) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        # output = th.where((t == 0), decoder_nll, kl)
+        output = kl
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs, analog_bit, noise=None):
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+        tmp_mask = (1 - model_kwargs['src_key_padding_mask'])
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                padding_mask = tmp_mask,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            xtalpha = _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape).permute([0,2,1])
+            epsalpha = _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape).permute([0,2,1])
+            model_output_dec, model_output_bin = model(x_t, self._scale_timesteps(t), xtalpha=xtalpha, epsalpha=epsalpha, **model_kwargs)
+            # model_output_dec = model(x_t, self._scale_timesteps(t), **model_kwargs)
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    padding_mask = tmp_mask,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
+                    x_start=x_start, x_t=x_t, t=t
+                )[0],
+                ModelMeanType.START_X: x_start,
+                ModelMeanType.EPSILON: noise,
+            }[self.model_mean_type]
+
+            if not analog_bit:
+                def dec2bin(xinp, bits):
+                    mask = 2 ** th.arange(bits - 1, -1, -1).to(xinp.device, xinp.dtype)
+                    return xinp.unsqueeze(-1).bitwise_and(mask).ne(0).float()
+                bin_target = x_start.detach()
+                bin_target = (bin_target/2 + 0.5) # -> [0,1]
+                bin_target = bin_target * 256 #-> [0, 256]
+                bin_target = dec2bin(bin_target.permute([0,2,1]).round().int(), 8) 
+                bin_target = bin_target.reshape([target.shape[0], target.shape[2], 16]).permute([0,2,1])
+                t_weights = (t<10).cuda().unsqueeze(1).unsqueeze(2)
+                t_weights = t_weights * (t_weights.shape[0]/max(1, t_weights.sum()))
+                bin_target[bin_target==0] = -1
+                assert model_output_bin.shape == bin_target.shape
+
+            assert model_output_dec.shape == target.shape == x_start.shape
+
+            if not analog_bit:
+                terms["mse_bin"] = mean_flat(((bin_target - model_output_bin) ** 2) * t_weights, tmp_mask)
+            terms["mse_dec"] = mean_flat(((target - model_output_dec) ** 2), tmp_mask)
+
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                if not analog_bit:
+                    terms["loss"] = terms["mse_dec"] + terms["mse_bin"]
+                else:
+                    terms["loss"] = terms["mse_dec"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)

house_diffusion/logger.py

@@ -0,0 +1,496 @@
+"""
+Logger copied from OpenAI baselines to avoid extra RL-based dependencies:
+https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py
+"""
+
+import os
+import sys
+import shutil
+import os.path as osp
+import json
+import time
+import datetime
+import tempfile
+import warnings
+from collections import defaultdict
+from contextlib import contextmanager
+
+DEBUG = 10
+INFO = 20
+WARN = 30
+ERROR = 40
+
+DISABLED = 50
+
+
+class KVWriter(object):
+    def writekvs(self, kvs):
+        raise NotImplementedError
+
+
+class SeqWriter(object):
+    def writeseq(self, seq):
+        raise NotImplementedError
+
+
+class HumanOutputFormat(KVWriter, SeqWriter):
+    def __init__(self, filename_or_file):
+        if isinstance(filename_or_file, str):
+            self.file = open(filename_or_file, "wt")
+            self.own_file = True
+        else:
+            assert hasattr(filename_or_file, "read"), (
+                "expected file or str, got %s" % filename_or_file
+            )
+            self.file = filename_or_file
+            self.own_file = False
+
+    def writekvs(self, kvs):
+        # Create strings for printing
+        key2str = {}
+        for (key, val) in sorted(kvs.items()):
+            if hasattr(val, "__float__"):
+                valstr = "%-8.3g" % val
+            else:
+                valstr = str(val)
+            key2str[self._truncate(key)] = self._truncate(valstr)
+
+        # Find max widths
+        if len(key2str) == 0:
+            print("WARNING: tried to write empty key-value dict")
+            return
+        else:
+            keywidth = max(map(len, key2str.keys()))
+            valwidth = max(map(len, key2str.values()))
+
+        # Write out the data
+        dashes = "-" * (keywidth + valwidth + 7)
+        lines = [dashes]
+        for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()):
+            lines.append(
+                "| %s%s | %s%s |"
+                % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val)))
+            )
+        lines.append(dashes)
+        self.file.write("\n".join(lines) + "\n")
+
+        # Flush the output to the file
+        self.file.flush()
+
+    def _truncate(self, s):
+        maxlen = 30
+        return s[: maxlen - 3] + "..." if len(s) > maxlen else s
+
+    def writeseq(self, seq):
+        seq = list(seq)
+        for (i, elem) in enumerate(seq):
+            self.file.write(elem)
+            if i < len(seq) - 1:  # add space unless this is the last one
+                self.file.write(" ")
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        if self.own_file:
+            self.file.close()
+
+
+class JSONOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "wt")
+
+    def writekvs(self, kvs):
+        for k, v in sorted(kvs.items()):
+            if hasattr(v, "dtype"):
+                kvs[k] = float(v)
+        self.file.write(json.dumps(kvs) + "\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class CSVOutputFormat(KVWriter):
+    def __init__(self, filename):
+        self.file = open(filename, "w+t")
+        self.keys = []
+        self.sep = ","
+
+    def writekvs(self, kvs):
+        # Add our current row to the history
+        extra_keys = list(kvs.keys() - self.keys)
+        extra_keys.sort()
+        if extra_keys:
+            self.keys.extend(extra_keys)
+            self.file.seek(0)
+            lines = self.file.readlines()
+            self.file.seek(0)
+            for (i, k) in enumerate(self.keys):
+                if i > 0:
+                    self.file.write(",")
+                self.file.write(k)
+            self.file.write("\n")
+            for line in lines[1:]:
+                self.file.write(line[:-1])
+                self.file.write(self.sep * len(extra_keys))
+                self.file.write("\n")
+        for (i, k) in enumerate(self.keys):
+            if i > 0:
+                self.file.write(",")
+            v = kvs.get(k)
+            if v is not None:
+                self.file.write(str(v))
+        self.file.write("\n")
+        self.file.flush()
+
+    def close(self):
+        self.file.close()
+
+
+class TensorBoardOutputFormat(KVWriter):
+    """
+    Dumps key/value pairs into TensorBoard's numeric format.
+    """
+
+    def __init__(self, dir):
+        os.makedirs(dir, exist_ok=True)
+        self.dir = dir
+        self.step = 1
+        prefix = "events"
+        path = osp.join(osp.abspath(dir), prefix)
+        import tensorflow as tf
+        from tensorflow.python import pywrap_tensorflow
+        from tensorflow.core.util import event_pb2
+        from tensorflow.python.util import compat
+
+        self.tf = tf
+        self.event_pb2 = event_pb2
+        self.pywrap_tensorflow = pywrap_tensorflow
+        self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path))
+
+    def writekvs(self, kvs):
+        def summary_val(k, v):
+            kwargs = {"tag": k, "simple_value": float(v)}
+            return self.tf.Summary.Value(**kwargs)
+
+        summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()])
+        event = self.event_pb2.Event(wall_time=time.time(), summary=summary)
+        event.step = (
+            self.step
+        )  # is there any reason why you'd want to specify the step?
+        self.writer.WriteEvent(event)
+        self.writer.Flush()
+        self.step += 1
+
+    def close(self):
+        if self.writer:
+            self.writer.Close()
+            self.writer = None
+
+
+def make_output_format(format, ev_dir, log_suffix=""):
+    os.makedirs(ev_dir, exist_ok=True)
+    if format == "stdout":
+        return HumanOutputFormat(sys.stdout)
+    elif format == "log":
+        return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix))
+    elif format == "json":
+        return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix))
+    elif format == "csv":
+        return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix))
+    elif format == "tensorboard":
+        return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix))
+    else:
+        raise ValueError("Unknown format specified: %s" % (format,))
+
+
+# ================================================================
+# API
+# ================================================================
+
+
+def logkv(key, val):
+    """
+    Log a value of some diagnostic
+    Call this once for each diagnostic quantity, each iteration
+    If called many times, last value will be used.
+    """
+    get_current().logkv(key, val)
+
+
+def logkv_mean(key, val):
+    """
+    The same as logkv(), but if called many times, values averaged.
+    """
+    get_current().logkv_mean(key, val)
+
+
+def logkvs(d):
+    """
+    Log a dictionary of key-value pairs
+    """
+    for (k, v) in d.items():
+        logkv(k, v)
+
+
+def dumpkvs():
+    """
+    Write all of the diagnostics from the current iteration
+    """
+    return get_current().dumpkvs()
+
+
+def getkvs():
+    return get_current().name2val
+
+
+def log(*args, level=INFO):
+    """
+    Write the sequence of args, with no separators, to the console and output files (if you've configured an output file).
+    """
+    get_current().log(*args, level=level)
+
+
+def debug(*args):
+    log(*args, level=DEBUG)
+
+
+def info(*args):
+    log(*args, level=INFO)
+
+
+def warn(*args):
+    log(*args, level=WARN)
+
+
+def error(*args):
+    log(*args, level=ERROR)
+
+
+def set_level(level):
+    """
+    Set logging threshold on current logger.
+    """
+    get_current().set_level(level)
+
+
+def set_comm(comm):
+    get_current().set_comm(comm)
+
+
+def get_dir():
+    """
+    Get directory that log files are being written to.
+    will be None if there is no output directory (i.e., if you didn't call start)
+    """
+    return get_current().get_dir()
+
+
+record_tabular = logkv
+dump_tabular = dumpkvs
+
+
+@contextmanager
+def profile_kv(scopename):
+    logkey = "wait_" + scopename
+    tstart = time.time()
+    try:
+        yield
+    finally:
+        get_current().name2val[logkey] += time.time() - tstart
+
+
+def profile(n):
+    """
+    Usage:
+    @profile("my_func")
+    def my_func(): code
+    """
+
+    def decorator_with_name(func):
+        def func_wrapper(*args, **kwargs):
+            with profile_kv(n):
+                return func(*args, **kwargs)
+
+        return func_wrapper
+
+    return decorator_with_name
+
+
+# ================================================================
+# Backend
+# ================================================================
+
+
+def get_current():
+    if Logger.CURRENT is None:
+        _configure_default_logger()
+
+    return Logger.CURRENT
+
+
+class Logger(object):
+    DEFAULT = None  # A logger with no output files. (See right below class definition)
+    # So that you can still log to the terminal without setting up any output files
+    CURRENT = None  # Current logger being used by the free functions above
+
+    def __init__(self, dir, output_formats, comm=None):
+        self.name2val = defaultdict(float)  # values this iteration
+        self.name2cnt = defaultdict(int)
+        self.level = INFO
+        self.dir = dir
+        self.output_formats = output_formats
+        self.comm = comm
+
+    # Logging API, forwarded
+    # ----------------------------------------
+    def logkv(self, key, val):
+        self.name2val[key] = val
+
+    def logkv_mean(self, key, val):
+        oldval, cnt = self.name2val[key], self.name2cnt[key]
+        self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1)
+        self.name2cnt[key] = cnt + 1
+
+    def dumpkvs(self):
+        if self.comm is None:
+            d = self.name2val
+        else:
+            d = mpi_weighted_mean(
+                self.comm,
+                {
+                    name: (val, self.name2cnt.get(name, 1))
+                    for (name, val) in self.name2val.items()
+                },
+            )
+            if self.comm.rank != 0:
+                d["dummy"] = 1  # so we don't get a warning about empty dict
+        out = d.copy()  # Return the dict for unit testing purposes
+        for fmt in self.output_formats:
+            if isinstance(fmt, KVWriter):
+                fmt.writekvs(d)
+        self.name2val.clear()
+        self.name2cnt.clear()
+        return out
+
+    def log(self, *args, level=INFO):
+        if self.level <= level:
+            self._do_log(args)
+
+    # Configuration
+    # ----------------------------------------
+    def set_level(self, level):
+        self.level = level
+
+    def set_comm(self, comm):
+        self.comm = comm
+
+    def get_dir(self):
+        return self.dir
+
+    def close(self):
+        for fmt in self.output_formats:
+            fmt.close()
+
+    # Misc
+    # ----------------------------------------
+    def _do_log(self, args):
+        for fmt in self.output_formats:
+            if isinstance(fmt, SeqWriter):
+                fmt.writeseq(map(str, args))
+
+
+def get_rank_without_mpi_import():
+    # check environment variables here instead of importing mpi4py
+    # to avoid calling MPI_Init() when this module is imported
+    for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]:
+        if varname in os.environ:
+            return int(os.environ[varname])
+    return 0
+
+
+def mpi_weighted_mean(comm, local_name2valcount):
+    """
+    Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110
+    Perform a weighted average over dicts that are each on a different node
+    Input: local_name2valcount: dict mapping key -> (value, count)
+    Returns: key -> mean
+    """
+    all_name2valcount = comm.gather(local_name2valcount)
+    if comm.rank == 0:
+        name2sum = defaultdict(float)
+        name2count = defaultdict(float)
+        for n2vc in all_name2valcount:
+            for (name, (val, count)) in n2vc.items():
+                try:
+                    val = float(val)
+                except ValueError:
+                    if comm.rank == 0:
+                        warnings.warn(
+                            "WARNING: tried to compute mean on non-float {}={}".format(
+                                name, val
+                            )
+                        )
+                else:
+                    name2sum[name] += val * count
+                    name2count[name] += count
+        return {name: name2sum[name] / name2count[name] for name in name2sum}
+    else:
+        return {}
+
+
+def configure(dir=None, format_strs=None, comm=None, log_suffix=""):
+    """
+    If comm is provided, average all numerical stats across that comm
+    """
+    if dir is None:
+        dir = os.getenv("OPENAI_LOGDIR")
+    if dir is None:
+        dir = osp.join(
+            # tempfile.gettempdir(),
+            'ckpts',
+            datetime.datetime.now().strftime("openai_%Y_%m_%d_%H_%M_%S_%f"),
+        )
+    assert isinstance(dir, str)
+    dir = os.path.expanduser(dir)
+    os.makedirs(os.path.expanduser(dir), exist_ok=True)
+
+    rank = get_rank_without_mpi_import()
+    if rank > 0:
+        log_suffix = log_suffix + "-rank%03i" % rank
+
+    if format_strs is None:
+        if rank == 0:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",")
+        else:
+            format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",")
+    format_strs = filter(None, format_strs)
+    output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs]
+
+    Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm)
+    if output_formats:
+        log("Logging to %s" % dir)
+
+
+def _configure_default_logger():
+    configure()
+    Logger.DEFAULT = Logger.CURRENT
+
+
+def reset():
+    if Logger.CURRENT is not Logger.DEFAULT:
+        Logger.CURRENT.close()
+        Logger.CURRENT = Logger.DEFAULT
+        log("Reset logger")
+
+
+@contextmanager
+def scoped_configure(dir=None, format_strs=None, comm=None):
+    prevlogger = Logger.CURRENT
+    configure(dir=dir, format_strs=format_strs, comm=comm)
+    try:
+        yield
+    finally:
+        Logger.CURRENT.close()
+        Logger.CURRENT = prevlogger
+

house_diffusion/losses.py

@@ -0,0 +1,77 @@
+"""
+Helpers for various likelihood-based losses. These are ported from the original
+Ho et al. diffusion models codebase:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py
+"""
+
+import numpy as np
+
+import torch as th
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs

house_diffusion/nn.py

@@ -0,0 +1,172 @@
+"""
+Various utilities for neural networks.
+"""
+
+import math
+
+import torch as th
+import torch.nn as nn
+
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+    def forward(self, x):
+        return x * th.sigmoid(x)
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def update_ema(target_params, source_params, rate=0.99):
+    """
+    Update target parameters to be closer to those of source parameters using
+    an exponential moving average.
+
+    :param target_params: the target parameter sequence.
+    :param source_params: the source parameter sequence.
+    :param rate: the EMA rate (closer to 1 means slower).
+    """
+    for targ, src in zip(target_params, source_params):
+        targ.detach().mul_(rate).add_(src, alpha=1 - rate)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def mean_flat(tensor, padding_mask):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    tensor = tensor * padding_mask.unsqueeze(1)
+    tensor = tensor.mean(dim=list(range(1, len(tensor.shape))))/th.sum(padding_mask, dim=1)
+    return tensor
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNorm32(32, channels)
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = th.exp(
+        -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = th.cat([th.cos(args), th.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return CheckpointFunction.apply(func, len(inputs), *args)
+    else:
+        return func(*inputs)
+
+
+class CheckpointFunction(th.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        with th.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with th.enable_grad():
+            # Fixes a bug where the first op in run_function modifies the
+            # Tensor storage in place, which is not allowed for detach()'d
+            # Tensors.
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = th.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        del ctx.input_tensors
+        del ctx.input_params
+        del output_tensors
+        return (None, None) + input_grads

house_diffusion/resample.py

@@ -0,0 +1,154 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch as th
+import torch.distributed as dist
+
+
+def create_named_schedule_sampler(name, diffusion):
+    """
+    Create a ScheduleSampler from a library of pre-defined samplers.
+
+    :param name: the name of the sampler.
+    :param diffusion: the diffusion object to sample for.
+    """
+    if name == "uniform":
+        return UniformSampler(diffusion)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(diffusion)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {name}")
+
+
+class ScheduleSampler(ABC):
+    """
+    A distribution over timesteps in the diffusion process, intended to reduce
+    variance of the objective.
+
+    By default, samplers perform unbiased importance sampling, in which the
+    objective's mean is unchanged.
+    However, subclasses may override sample() to change how the resampled
+    terms are reweighted, allowing for actual changes in the objective.
+    """
+
+    @abstractmethod
+    def weights(self):
+        """
+        Get a numpy array of weights, one per diffusion step.
+
+        The weights needn't be normalized, but must be positive.
+        """
+
+    def sample(self, batch_size, device):
+        """
+        Importance-sample timesteps for a batch.
+
+        :param batch_size: the number of timesteps.
+        :param device: the torch device to save to.
+        :return: a tuple (timesteps, weights):
+                 - timesteps: a tensor of timestep indices.
+                 - weights: a tensor of weights to scale the resulting losses.
+        """
+        w = self.weights()
+        p = w / np.sum(w)
+        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
+        indices = th.from_numpy(indices_np).long().to(device)
+        weights_np = 1 / (len(p) * p[indices_np])
+        weights = th.from_numpy(weights_np).float().to(device)
+        return indices, weights
+
+
+class UniformSampler(ScheduleSampler):
+    def __init__(self, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()

house_diffusion/respace.py

@@ -0,0 +1,128 @@
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(
+            model, self.timestep_map, self.rescale_timesteps, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        if self.rescale_timesteps:
+            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)

house_diffusion/rplanhg_datasets.py

@@ -0,0 +1,620 @@
+import math
+import random
+import torch as th
+
+from PIL import Image, ImageDraw
+import blobfile as bf
+from mpi4py import MPI
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+from glob import glob
+import json
+import os
+import cv2 as cv
+from tqdm import tqdm
+from shapely import geometry as gm
+from shapely.ops import unary_union
+from collections import defaultdict
+import copy
+
+
+def load_rplanhg_data(
+        batch_size,
+        analog_bit,
+        target_set,
+        set_name='train',
+):
+    """
+    For a dataset, create a generator over (shapes, kwargs) pairs.
+    """
+    # set_name = 'train'
+    set_name = 'eval'
+    print(f"loading {set_name} of target set {target_set}")
+    deterministic = False if set_name == 'train' else True
+    dataset = RPlanhgDataset(set_name, analog_bit, target_set)
+    if deterministic:
+        loader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=False, num_workers=2, drop_last=False
+        )
+    else:
+        loader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=False
+        )
+    while True:
+        yield from loader
+
+
+def make_non_manhattan(poly, polygon, house_poly):
+    dist = abs(poly[2] - poly[0])
+    direction = np.argmin(dist)
+    center = poly.mean(0)
+    min = poly.min(0)
+    max = poly.max(0)
+
+    tmp = np.random.randint(3, 7)
+    new_min_y = center[1] - (max[1] - min[1]) / tmp
+    new_max_y = center[1] + (max[1] - min[1]) / tmp
+    if center[0] < 128:
+        new_min_x = min[0] - (max[0] - min[0]) / np.random.randint(2, 5)
+        new_max_x = center[0]
+        poly1 = [[min[0], min[1]], [new_min_x, new_min_y], [new_min_x, new_max_y], [min[0], max[1]], [max[0], max[1]],
+                 [max[0], min[1]]]
+    else:
+        new_min_x = center[0]
+        new_max_x = max[0] + (max[0] - min[0]) / np.random.randint(2, 5)
+        poly1 = [[min[0], min[1]], [min[0], max[1]], [max[0], max[1]], [new_max_x, new_max_y], [new_max_x, new_min_y],
+                 [max[0], min[1]]]
+
+    new_min_x = center[0] - (max[0] - min[0]) / tmp
+    new_max_x = center[0] + (max[0] - min[0]) / tmp
+    if center[1] < 128:
+        new_min_y = min[1] - (max[1] - min[1]) / np.random.randint(2, 5)
+        new_max_y = center[1]
+        poly2 = [[min[0], min[1]], [min[0], max[1]], [max[0], max[1]], [max[0], min[1]], [new_max_x, new_min_y],
+                 [new_min_x, new_min_y]]
+    else:
+        new_min_y = center[1]
+        new_max_y = max[1] + (max[1] - min[1]) / np.random.randint(2, 5)
+        poly2 = [[min[0], min[1]], [min[0], max[1]], [new_min_x, new_max_y], [new_max_x, new_max_y], [max[0], max[1]],
+                 [max[0], min[1]]]
+    p1 = gm.Polygon(poly1)
+    iou1 = house_poly.intersection(p1).area / p1.area
+    p2 = gm.Polygon(poly2)
+    iou2 = house_poly.intersection(p2).area / p2.area
+    if iou1 > 0.9 and iou2 > 0.9:
+        return poly
+    if iou1 < iou2:
+        return poly1
+    else:
+        return poly2
+
+
+get_bin = lambda x, z: [int(y) for y in format(x, 'b').zfill(z)]
+get_one_hot = lambda x, z: np.eye(z)[min(x, z - 1)]
+
+
+class RPlanhgDataset(Dataset):
+    def __init__(self, set_name, analog_bit, target_set, non_manhattan=False):
+        super().__init__()
+        base_dir = '../datasets/rplan'
+        self.non_manhattan = non_manhattan
+        self.set_name = set_name
+        self.analog_bit = analog_bit
+        self.target_set = target_set
+        self.subgraphs = []
+        self.org_graphs = []
+        self.org_houses = []
+        max_num_points = 100
+        if self.set_name == 'eval':
+            cnumber_dist = np.load(f'processed_rplan/rplan_train_{target_set}_cndist.npz', allow_pickle=True)[
+                'cnumber_dist'].item()
+        if os.path.exists(f'processed_rplan/rplan_{set_name}_{target_set}.npz'):
+            data = np.load(f'processed_rplan/rplan_{set_name}_{target_set}.npz', allow_pickle=True)
+            self.graphs = data['graphs']
+            self.houses = data['houses']
+            self.door_masks = data['door_masks']
+            self.self_masks = data['self_masks']
+            self.gen_masks = data['gen_masks']
+            self.num_coords = 2
+            self.max_num_points = max_num_points
+            cnumber_dist = np.load(f'processed_rplan/rplan_train_{target_set}_cndist.npz', allow_pickle=True)[
+                'cnumber_dist'].item()
+            if self.set_name == 'eval':
+                data = np.load(f'processed_rplan/rplan_{set_name}_{target_set}_syn.npz', allow_pickle=True)
+                self.syn_graphs = data['graphs']
+                self.syn_houses = data['houses']
+                self.syn_door_masks = data['door_masks']
+                self.syn_self_masks = data['self_masks']
+                self.syn_gen_masks = data['gen_masks']
+        else:
+            with open(f'{base_dir}/list.txt') as f:
+                lines = f.readlines()
+            cnt = 0
+
+            # TODO
+            failed_plans = []
+
+            for line in tqdm(lines):
+                #     cnt=cnt+1
+                #     file_name = f'{base_dir}/{line[:-1]}'
+                #     rms_type, fp_eds,rms_bbs,eds_to_rms=reader(file_name)
+                #     fp_size = len([x for x in rms_type if x != 15 and x != 17])
+                #     if self.set_name=='train' and fp_size == target_set:
+                #             continue
+                #     if self.set_name=='eval' and fp_size != target_set:
+                #             continue
+                #     a = [rms_type, rms_bbs, fp_eds, eds_to_rms]
+                #     self.subgraphs.append(a)
+
+                # for graph in tqdm(self.subgraphs):
+                try:
+                    cnt = cnt + 1
+                    file_name = f'{base_dir}/{line[:-1]}'
+                    rms_type, fp_eds, rms_bbs, eds_to_rms = reader(file_name)
+                    fp_size = len([x for x in rms_type if x != 15 and x != 17])
+                    if self.set_name == 'train' and fp_size == target_set:
+                        continue
+                    if self.set_name == 'eval' and fp_size != target_set:
+                        continue
+                    graph = [rms_type, rms_bbs, fp_eds, eds_to_rms]
+                    rms_type = graph[0]
+                    rms_bbs = graph[1]
+                    fp_eds = graph[2]
+                    eds_to_rms = graph[3]
+                    rms_bbs = np.array(rms_bbs)
+                    fp_eds = np.array(fp_eds)
+
+                    # extract boundary box and centralize
+                    tl = np.min(rms_bbs[:, :2], 0)
+                    br = np.max(rms_bbs[:, 2:], 0)
+                    shift = (tl + br) / 2.0 - 0.5
+                    rms_bbs[:, :2] -= shift
+                    rms_bbs[:, 2:] -= shift
+                    fp_eds[:, :2] -= shift
+                    fp_eds[:, 2:] -= shift
+                    tl -= shift
+                    br -= shift
+
+                    # build input graph
+                    graph_nodes, graph_edges, rooms_mks = self.build_graph(rms_type, fp_eds, eds_to_rms)
+
+                    house = []
+                    for room_mask, room_type in zip(rooms_mks, graph_nodes):
+                        room_mask = room_mask.astype(np.uint8)
+                        room_mask = cv.resize(room_mask, (256, 256), interpolation=cv.INTER_AREA)
+                        contours, _ = cv.findContours(room_mask, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
+                        contours = contours[0]
+                        house.append([contours[:, 0, :], room_type])
+                    self.org_graphs.append(graph_edges)
+                    self.org_houses.append(house)
+
+                except IndexError:
+                    # print(line)
+                    failed_plans.append(line)
+            print("failed: ", failed_plans)
+            print("len: ", len(failed_plans))
+
+            houses = []
+            door_masks = []
+            self_masks = []
+            gen_masks = []
+            graphs = []
+            if self.set_name == 'train':
+                cnumber_dist = defaultdict(list)
+
+            if self.non_manhattan:
+                for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
+                    # Generating non-manhattan Balconies
+                    tmp = []
+                    for i, room in enumerate(h):
+                        if room[1] > 10:
+                            continue
+                        if len(room[0]) != 4:
+                            continue
+                        if np.random.randint(2):
+                            continue
+                        poly = gm.Polygon(room[0])
+                        house_polygon = unary_union([gm.Polygon(room[0]) for room in h])
+                        room[0] = make_non_manhattan(room[0], poly, house_polygon)
+
+            for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
+                house = []
+                corner_bounds = []
+                num_points = 0
+                for i, room in enumerate(h):
+                    if room[1] > 10:
+                        room[1] = {15: 11, 17: 12, 16: 13}[room[1]]
+                    room[0] = np.reshape(room[0], [len(room[0]),
+                                                   2]) / 256. - 0.5  # [[x0,y0],[x1,y1],...,[x15,y15]] and map to 0-1 - > -0.5, 0.5
+                    room[0] = room[0] * 2  # map to [-1, 1]
+                    if self.set_name == 'train':
+                        cnumber_dist[room[1]].append(len(room[0]))
+                    # Adding conditions
+                    num_room_corners = len(room[0])
+                    rtype = np.repeat(np.array([get_one_hot(room[1], 25)]), num_room_corners, 0)
+                    room_index = np.repeat(np.array([get_one_hot(len(house) + 1, 32)]), num_room_corners, 0)
+                    corner_index = np.array([get_one_hot(x, 32) for x in range(num_room_corners)])
+                    # Src_key_padding_mask
+                    padding_mask = np.repeat(1, num_room_corners)
+                    padding_mask = np.expand_dims(padding_mask, 1)
+                    # Generating corner bounds for attention masks
+                    connections = np.array([[i, (i + 1) % num_room_corners] for i in range(num_room_corners)])
+                    connections += num_points
+                    corner_bounds.append([num_points, num_points + num_room_corners])
+                    num_points += num_room_corners
+                    room = np.concatenate((room[0], rtype, corner_index, room_index, padding_mask, connections), 1)
+                    house.append(room)
+
+                house_layouts = np.concatenate(house, 0)
+                if len(house_layouts) > max_num_points:
+                    continue
+                padding = np.zeros((max_num_points - len(house_layouts), 94))
+                gen_mask = np.ones((max_num_points, max_num_points))
+                gen_mask[:len(house_layouts), :len(house_layouts)] = 0
+                house_layouts = np.concatenate((house_layouts, padding), 0)
+
+                door_mask = np.ones((max_num_points, max_num_points))
+                self_mask = np.ones((max_num_points, max_num_points))
+                for i in range(len(corner_bounds)):
+                    for j in range(len(corner_bounds)):
+                        if i == j:
+                            self_mask[corner_bounds[i][0]:corner_bounds[i][1],
+                            corner_bounds[j][0]:corner_bounds[j][1]] = 0
+                        elif any(np.equal([i, 1, j], graph).all(1)) or any(np.equal([j, 1, i], graph).all(1)):
+                            door_mask[corner_bounds[i][0]:corner_bounds[i][1],
+                            corner_bounds[j][0]:corner_bounds[j][1]] = 0
+                houses.append(house_layouts)
+                door_masks.append(door_mask)
+                self_masks.append(self_mask)
+                gen_masks.append(gen_mask)
+                graphs.append(graph)
+            self.max_num_points = max_num_points
+            self.houses = houses
+            self.door_masks = door_masks
+            self.self_masks = self_masks
+            self.gen_masks = gen_masks
+            self.num_coords = 2
+            self.graphs = graphs
+
+            # --------------
+            # graph_dict = {f'graph_{i}': graph for i, graph in enumerate(self.graphs)}
+            for i, graph in enumerate(self.graphs):
+                print(f"Graph {i}: shape = {np.shape(graph)}, type = {type(graph)}")
+
+            # Save each graph individually within a dictionary
+            # graph_dict = {f'graph_{i}': graph for i, graph in enumerate(self.graphs)}
+
+            np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}', graphs=self.graphs,
+                                houses=self.houses,
+                                # np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}', **graph_dict, houses=self.houses,
+                                door_masks=self.door_masks, self_masks=self.self_masks, gen_masks=self.gen_masks)
+            if self.set_name == 'train':
+                np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}_cndist', cnumber_dist=cnumber_dist)
+
+            if set_name == 'eval':
+                houses = []
+                graphs = []
+                door_masks = []
+                self_masks = []
+                gen_masks = []
+                len_house_layouts = 0
+                for h, graph in tqdm(zip(self.org_houses, self.org_graphs), desc='processing dataset'):
+                    house = []
+                    corner_bounds = []
+                    num_points = 0
+                    # num_room_corners_total = [cnumber_dist[room[1]][random.randint(0, len(cnumber_dist[room[1]])-1)] for room in h]
+                    # while np.sum(num_room_corners_total)>=max_num_points:
+                    #     num_room_corners_total = [cnumber_dist[room[1]][random.randint(0, len(cnumber_dist[room[1]])-1)] for room in h]
+                    num_room_corners_total = []
+                    for room in h:
+                        room_type = room[1]
+                        default_value = 4
+                        if room_type in cnumber_dist and cnumber_dist[room_type]:
+                            num_room_corners_total.append(
+                                cnumber_dist[room_type][random.randint(0, len(cnumber_dist[room_type]) - 1)]
+                            )
+                        else:
+                            # Handle the case where cnumber_dist[room_type] is missing or empty
+                            print(f"Warning: No data found for room type {room_type}. Assigning default value.")
+                            default_value = 4  # Assign a reasonable default value or handle accordingly
+                            num_room_corners_total.append(default_value)
+
+                    while np.sum(num_room_corners_total) >= max_num_points:
+                        num_room_corners_total = []
+                        for room in h:
+                            room_type = room[1]
+                            if room_type in cnumber_dist and cnumber_dist[room_type]:
+                                num_room_corners_total.append(
+                                    cnumber_dist[room_type][random.randint(0, len(cnumber_dist[room_type]) - 1)]
+                                )
+                        else:
+                            num_room_corners_total.append(default_value)
+
+                    for i, room in enumerate(h):
+                        # Adding conditions
+                        num_room_corners = num_room_corners_total[i]
+                        rtype = np.repeat(np.array([get_one_hot(room[1], 25)]), num_room_corners, 0)
+                        room_index = np.repeat(np.array([get_one_hot(len(house) + 1, 32)]), num_room_corners, 0)
+                        corner_index = np.array([get_one_hot(x, 32) for x in range(num_room_corners)])
+                        # Src_key_padding_mask
+                        padding_mask = np.repeat(1, num_room_corners)
+                        padding_mask = np.expand_dims(padding_mask, 1)
+                        # Generating corner bounds for attention masks
+                        connections = np.array([[i, (i + 1) % num_room_corners] for i in range(num_room_corners)])
+                        connections += num_points
+                        corner_bounds.append([num_points, num_points + num_room_corners])
+                        num_points += num_room_corners
+                        room = np.concatenate((np.zeros([num_room_corners, 2]), rtype, corner_index, room_index,
+                                               padding_mask, connections), 1)
+                        house.append(room)
+
+                    house_layouts = np.concatenate(house, 0)
+                    if np.sum([len(room[0]) for room in h]) > max_num_points:
+                        continue
+                    padding = np.zeros((max_num_points - len(house_layouts), 94))
+                    gen_mask = np.ones((max_num_points, max_num_points))
+                    gen_mask[:len(house_layouts), :len(house_layouts)] = 0
+                    house_layouts = np.concatenate((house_layouts, padding), 0)
+
+                    door_mask = np.ones((max_num_points, max_num_points))
+                    self_mask = np.ones((max_num_points, max_num_points))
+                    for i, room in enumerate(h):
+                        if room[1] == 1:
+                            living_room_index = i
+                            break
+                    for i in range(len(corner_bounds)):
+                        is_connected = False
+                        for j in range(len(corner_bounds)):
+                            if i == j:
+                                self_mask[corner_bounds[i][0]:corner_bounds[i][1],
+                                corner_bounds[j][0]:corner_bounds[j][1]] = 0
+                            elif any(np.equal([i, 1, j], graph).all(1)) or any(np.equal([j, 1, i], graph).all(1)):
+                                door_mask[corner_bounds[i][0]:corner_bounds[i][1],
+                                corner_bounds[j][0]:corner_bounds[j][1]] = 0
+                                is_connected = True
+                        if not is_connected:
+                            door_mask[corner_bounds[i][0]:corner_bounds[i][1],
+                            corner_bounds[living_room_index][0]:corner_bounds[living_room_index][1]] = 0
+
+                    houses.append(house_layouts)
+                    door_masks.append(door_mask)
+                    self_masks.append(self_mask)
+                    gen_masks.append(gen_mask)
+                    graphs.append(graph)
+                self.syn_houses = houses
+                self.syn_door_masks = door_masks
+                self.syn_self_masks = self_masks
+                self.syn_gen_masks = gen_masks
+                self.syn_graphs = graphs
+                np.savez_compressed(f'processed_rplan/rplan_{set_name}_{target_set}_syn', graphs=self.syn_graphs,
+                                    houses=self.syn_houses,
+                                    door_masks=self.syn_door_masks, self_masks=self.syn_self_masks,
+                                    gen_masks=self.syn_gen_masks)
+
+
+    def __len__(self):
+        return len(self.houses)
+
+    def __getitem__(self, idx):
+        # idx = int(idx//20)
+        arr = self.houses[idx][:, :self.num_coords]
+        graph = np.concatenate((self.graphs[idx], np.zeros([200 - len(self.graphs[idx]), 3])), 0)
+
+        cond = {
+            'door_mask': self.door_masks[idx],
+            'self_mask': self.self_masks[idx],
+            'gen_mask': self.gen_masks[idx],
+            'room_types': self.houses[idx][:, self.num_coords:self.num_coords + 25],
+            'corner_indices': self.houses[idx][:, self.num_coords + 25:self.num_coords + 57],
+            'room_indices': self.houses[idx][:, self.num_coords + 57:self.num_coords + 89],
+            'src_key_padding_mask': 1 - self.houses[idx][:, self.num_coords + 89],
+            'connections': self.houses[idx][:, self.num_coords + 90:self.num_coords + 92],
+            'graph': graph,
+        }
+        if self.set_name == 'eval':
+            syn_graph = np.concatenate((self.syn_graphs[idx], np.zeros([200 - len(self.syn_graphs[idx]), 3])), 0)
+            assert (graph == syn_graph).all(), idx
+            cond.update({
+                'syn_door_mask': self.syn_door_masks[idx],
+                'syn_self_mask': self.syn_self_masks[idx],
+                'syn_gen_mask': self.syn_gen_masks[idx],
+                'syn_room_types': self.syn_houses[idx][:, self.num_coords:self.num_coords + 25],
+                'syn_corner_indices': self.syn_houses[idx][:, self.num_coords + 25:self.num_coords + 57],
+                'syn_room_indices': self.syn_houses[idx][:, self.num_coords + 57:self.num_coords + 89],
+                'syn_src_key_padding_mask': 1 - self.syn_houses[idx][:, self.num_coords + 89],
+                'syn_connections': self.syn_houses[idx][:, self.num_coords + 90:self.num_coords + 92],
+                'syn_graph': syn_graph,
+            })
+        if self.set_name == 'train':
+            #### Random Rotate
+            rotation = random.randint(0, 3)
+            if rotation == 1:
+                arr[:, [0, 1]] = arr[:, [1, 0]]
+                arr[:, 0] = -arr[:, 0]
+            elif rotation == 2:
+                arr[:, [0, 1]] = -arr[:, [1, 0]]
+            elif rotation == 3:
+                arr[:, [0, 1]] = arr[:, [1, 0]]
+                arr[:, 1] = -arr[:, 1]
+
+            ## To generate any rotation uncomment this
+
+            # if self.non_manhattan:
+            # theta = random.random()*np.pi/2
+            # rot_mat = np.array([[np.cos(theta), -np.sin(theta), 0],
+            # [np.sin(theta), np.cos(theta), 0]])
+            # arr = np.matmul(arr,rot_mat)[:,:2]
+
+            # Random Scale
+            # arr = arr * np.random.normal(1., .5)
+
+            # Random Shift
+            # arr[:, 0] = arr[:, 0] + np.random.normal(0., .1)
+            # arr[:, 1] = arr[:, 1] + np.random.normal(0., .1)
+
+        if not self.analog_bit:
+            arr = np.transpose(arr, [1, 0])
+            return arr.astype(float), cond
+        else:
+            ONE_HOT_RES = 256
+            arr_onehot = np.zeros((ONE_HOT_RES * 2, arr.shape[1])) - 1
+            xs = ((arr[:, 0] + 1) * (ONE_HOT_RES / 2)).astype(int)
+            ys = ((arr[:, 1] + 1) * (ONE_HOT_RES / 2)).astype(int)
+            xs = np.array([get_bin(x, 8) for x in xs])
+            ys = np.array([get_bin(x, 8) for x in ys])
+            arr_onehot = np.concatenate([xs, ys], 1)
+            arr_onehot = np.transpose(arr_onehot, [1, 0])
+            arr_onehot[arr_onehot == 0] = -1
+            return arr_onehot.astype(float), cond
+
+    def make_sequence(self, edges):
+        polys = []
+        v_curr = tuple(edges[0][:2])
+        e_ind_curr = 0
+        e_visited = [0]
+        seq_tracker = [v_curr]
+        find_next = False
+        while len(e_visited) < len(edges):
+            if find_next == False:
+                if v_curr == tuple(edges[e_ind_curr][2:]):
+                    v_curr = tuple(edges[e_ind_curr][:2])
+                else:
+                    v_curr = tuple(edges[e_ind_curr][2:])
+                find_next = not find_next
+            else:
+                # look for next edge
+                for k, e in enumerate(edges):
+                    if k not in e_visited:
+                        if (v_curr == tuple(e[:2])):
+                            v_curr = tuple(e[2:])
+                            e_ind_curr = k
+                            e_visited.append(k)
+                            break
+                        elif (v_curr == tuple(e[2:])):
+                            v_curr = tuple(e[:2])
+                            e_ind_curr = k
+                            e_visited.append(k)
+                            break
+
+            # extract next sequence
+            if v_curr == seq_tracker[-1]:
+                polys.append(seq_tracker)
+                for k, e in enumerate(edges):
+                    if k not in e_visited:
+                        v_curr = tuple(edges[0][:2])
+                        seq_tracker = [v_curr]
+                        find_next = False
+                        e_ind_curr = k
+                        e_visited.append(k)
+                        break
+            else:
+                seq_tracker.append(v_curr)
+        polys.append(seq_tracker)
+
+        return polys
+
+    def build_graph(self, rms_type, fp_eds, eds_to_rms, out_size=64):
+        # create edges
+        triples = []
+        nodes = rms_type
+        # encode connections
+        for k in range(len(nodes)):
+            for l in range(len(nodes)):
+                if l > k:
+                    is_adjacent = any([True for e_map in eds_to_rms if (l in e_map) and (k in e_map)])
+                    if is_adjacent:
+                        if 'train' in self.set_name:
+                            triples.append([k, 1, l])
+                        else:
+                            triples.append([k, 1, l])
+                    else:
+                        if 'train' in self.set_name:
+                            triples.append([k, -1, l])
+                        else:
+                            triples.append([k, -1, l])
+        # get rooms masks
+        eds_to_rms_tmp = []
+        for l in range(len(eds_to_rms)):
+            eds_to_rms_tmp.append([eds_to_rms[l][0]])
+        rms_masks = []
+        im_size = 256
+        fp_mk = np.zeros((out_size, out_size))
+        for k in range(len(nodes)):
+            # add rooms and doors
+            eds = []
+            for l, e_map in enumerate(eds_to_rms_tmp):
+                if (k in e_map):
+                    eds.append(l)
+            # draw rooms
+            rm_im = Image.new('L', (im_size, im_size))
+            dr = ImageDraw.Draw(rm_im)
+            for eds_poly in [eds]:
+                poly = self.make_sequence(np.array([fp_eds[l][:4] for l in eds_poly]))[0]
+                poly = [(im_size * x, im_size * y) for x, y in poly]
+                if len(poly) >= 2:
+                    dr.polygon(poly, fill='white')
+                else:
+                    print("Empty room")
+                    exit(0)
+            rm_im = rm_im.resize((out_size, out_size))
+            rm_arr = np.array(rm_im)
+            inds = np.where(rm_arr > 0)
+            rm_arr[inds] = 1.0
+            rms_masks.append(rm_arr)
+            if rms_type[k] != 15 and rms_type[k] != 17:
+                fp_mk[inds] = k + 1
+        # trick to remove overlap
+        for k in range(len(nodes)):
+            if rms_type[k] != 15 and rms_type[k] != 17:
+                rm_arr = np.zeros((out_size, out_size))
+                inds = np.where(fp_mk == k + 1)
+                rm_arr[inds] = 1.0
+                rms_masks[k] = rm_arr
+        # convert to array
+        nodes = np.array(nodes)
+        triples = np.array(triples)
+        rms_masks = np.array(rms_masks)
+        return nodes, triples, rms_masks
+
+
+def is_adjacent(box_a, box_b, threshold=0.03):
+    x0, y0, x1, y1 = box_a
+    x2, y2, x3, y3 = box_b
+    h1, h2 = x1 - x0, x3 - x2
+    w1, w2 = y1 - y0, y3 - y2
+    xc1, xc2 = (x0 + x1) / 2.0, (x2 + x3) / 2.0
+    yc1, yc2 = (y0 + y1) / 2.0, (y2 + y3) / 2.0
+    delta_x = np.abs(xc2 - xc1) - (h1 + h2) / 2.0
+    delta_y = np.abs(yc2 - yc1) - (w1 + w2) / 2.0
+    delta = max(delta_x, delta_y)
+    return delta < threshold
+
+
+def reader(filename):
+    with open(filename) as f:
+        info = json.load(f)
+        rms_bbs = np.asarray(info['boxes'])
+        fp_eds = info['edges']
+        rms_type = info['room_type']
+        eds_to_rms = info['ed_rm']
+        s_r = 0
+        for rmk in range(len(rms_type)):
+            if (rms_type[rmk] != 17):
+                s_r = s_r + 1
+        rms_bbs = np.array(rms_bbs) / 256.0
+        fp_eds = np.array(fp_eds) / 256.0
+        fp_eds = fp_eds[:, :4]
+        tl = np.min(rms_bbs[:, :2], 0)
+        br = np.max(rms_bbs[:, 2:], 0)
+        shift = (tl + br) / 2.0 - 0.5
+        rms_bbs[:, :2] -= shift
+        rms_bbs[:, 2:] -= shift
+        fp_eds[:, :2] -= shift
+        fp_eds[:, 2:] -= shift
+        tl -= shift
+        br -= shift
+        return rms_type, fp_eds, rms_bbs, eds_to_rms
+
+
+if __name__ == '__main__':
+    dataset = RPlanhgDataset('eval', False, 8)

house_diffusion/script_util.py

@@ -0,0 +1,182 @@
+import argparse
+import inspect
+
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+from .transformer import TransformerModel
+
+def diffusion_defaults():
+    """
+    Defaults for image and classifier training.
+    """
+    return dict(
+        analog_bit=False,
+        learn_sigma=False,
+        # diffusion_steps=25,
+        diffusion_steps=1000,
+        noise_schedule="cosine",
+        timestep_respacing="ddim100",
+        use_kl=False,
+        predict_xstart=False,
+        rescale_timesteps=False,
+        rescale_learned_sigmas=False,
+        # target_set=-1,
+
+        # target_set=4,
+        # target_set=5,
+        # target_set=6,
+        # target_set=7,
+        target_set=8,
+
+        set_name='',
+    )
+
+def update_arg_parser(args):
+    args.num_channels = 512
+    num_coords = 16 if args.analog_bit else 2
+    if args.dataset=='rplan':
+        args.input_channels = num_coords + (2*8 if not args.analog_bit else 0) # . , . , . , . , '
+        args.condition_channels = 89
+        args.out_channels = num_coords * 1
+        args.use_unet = False
+
+    elif args.dataset=='st3d':
+        args.input_channels = num_coords + (2*8 if not args.analog_bit else 0) # . , . , . , . , '
+        args.condition_channels = 89
+        args.out_channels = num_coords * 1
+        args.use_unet = False
+
+    elif args.dataset=='zind':
+        args.input_channels = num_coords + 2 * 8
+        args.condition_channels = 89
+        args.out_channels = num_coords * 1
+        args.use_unet = False
+
+    elif args.dataset=='layout':
+        args.use_unet = True
+        pass #TODO NEED TO COMPLETE
+
+    elif args.dataset=='outdoor':
+        args.use_unet = True
+        pass #TODO NEED TO COMPLETE
+    else:
+        assert False, "DATASET NOT FOUND"
+
+def model_and_diffusion_defaults():
+    """
+    Defaults for image training.
+    """
+    res = dict(
+            dataset='rplan',
+            # dataset='',
+            use_checkpoint=False,
+            input_channels=0,
+            condition_channels=0,
+            out_channels=0,
+            use_unet=False,
+            num_channels=128
+            )
+    res.update(diffusion_defaults())
+    return res
+
+def create_model_and_diffusion(
+    input_channels,
+    condition_channels,
+    num_channels,
+    out_channels,
+    dataset,
+    use_checkpoint,
+    use_unet,
+    learn_sigma,
+    diffusion_steps,
+    noise_schedule,
+    timestep_respacing,
+    use_kl,
+    predict_xstart,
+    rescale_timesteps,
+    rescale_learned_sigmas,
+    analog_bit,
+    target_set,
+    set_name,
+):
+    model = TransformerModel(input_channels, condition_channels, num_channels, out_channels, dataset, use_checkpoint, use_unet, analog_bit)
+
+    diffusion = create_gaussian_diffusion(
+        steps=diffusion_steps,
+        learn_sigma=learn_sigma,
+        noise_schedule=noise_schedule,
+        use_kl=use_kl,
+        predict_xstart=predict_xstart,
+        rescale_timesteps=rescale_timesteps,
+        rescale_learned_sigmas=rescale_learned_sigmas,
+        timestep_respacing=timestep_respacing,
+    )
+    return model, diffusion
+
+def create_gaussian_diffusion(
+    *,
+    steps=1000,
+    learn_sigma=False,
+    sigma_small=False,
+    noise_schedule="linear",
+    use_kl=False,
+    predict_xstart=False,
+    rescale_timesteps=False,
+    rescale_learned_sigmas=False,
+    timestep_respacing="",
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, steps)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if not timestep_respacing:
+        timestep_respacing = [steps]
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=(
+            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
+        ),
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type,
+        rescale_timesteps=rescale_timesteps,
+    )
+
+
+def add_dict_to_argparser(parser, default_dict):
+    for k, v in default_dict.items():
+        v_type = type(v)
+        if v is None:
+            v_type = str
+        elif isinstance(v, bool):
+            v_type = str2bool
+        parser.add_argument(f"--{k}", default=v, type=v_type)
+
+
+def args_to_dict(args, keys):
+    return {k: getattr(args, k) for k in keys}
+
+
+def str2bool(v):
+    """
+    https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
+    """
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("boolean value expected")

house_diffusion/train_util.py

@@ -0,0 +1,416 @@
+import copy
+import functools
+import os
+
+import blobfile as bf
+import torch as th
+import torch.distributed as dist
+from torch.nn.parallel.distributed import DistributedDataParallel as DDP
+from torch.optim import AdamW
+
+from . import dist_util, logger
+from .fp16_util import MixedPrecisionTrainer
+from .nn import update_ema
+from .resample import LossAwareSampler, UniformSampler
+
+# For ImageNet experiments, this was a good default value.
+# We found that the lg_loss_scale quickly climbed to
+# 20-21 within the first ~1K steps of training.
+INITIAL_LOG_LOSS_SCALE = 20.0
+
+
+class TrainLoop:
+    def __init__(
+            self,
+            *,
+            model,
+            diffusion,
+            data,
+            batch_size,
+            microbatch,
+            lr,
+            ema_rate,
+            log_interval,
+            save_interval,
+            resume_checkpoint,
+            use_fp16=False,
+            fp16_scale_growth=1e-3,
+            schedule_sampler=None,
+            weight_decay=0.0,
+            lr_anneal_steps=0,
+            analog_bit=None,
+    ):
+        self.analog_bit = analog_bit
+        self.model = model
+        self.diffusion = diffusion
+        self.data = data
+        self.batch_size = batch_size
+        self.microbatch = microbatch if microbatch > 0 else batch_size
+        self.lr = lr
+        self.ema_rate = (
+            [ema_rate]
+            if isinstance(ema_rate, float)
+            else [float(x) for x in ema_rate.split(",")]
+        )
+        self.log_interval = log_interval
+        self.save_interval = save_interval
+        self.resume_checkpoint = resume_checkpoint
+        self.use_fp16 = use_fp16
+        self.fp16_scale_growth = fp16_scale_growth
+        self.schedule_sampler = schedule_sampler or UniformSampler(diffusion)
+        self.weight_decay = weight_decay
+        self.lr_anneal_steps = lr_anneal_steps
+
+        self.step = 0
+        self.resume_step = 0
+        self.global_batch = self.batch_size * dist.get_world_size()
+
+        self.sync_cuda = th.cuda.is_available()
+
+        # TODO ------------------------------------------------------------------------
+        pretrained_path = "../ckpts/exp/model250000.pt"
+        pretrained_path = False
+
+        if pretrained_path:
+            self.load_pretrained(pretrained_path)
+        self.count_parameters_by_layer()
+
+        from .transformer_models import TransformerModels
+
+        device = th.device('cuda' if th.cuda.is_available() else 'cpu')
+        # self.model.to(device)
+        # print(th.get_default_device())
+        # th.set_default_device('cuda')
+        # print(th.get_default_device())
+
+        transformer_model = TransformerModels(self.model, device)
+        self.model_name = "Def"
+
+        # self.model = transformer_model.replace_InstanceNorm1d_LayerNorm()
+        # self.model_name = "Norm_LayerNorm"
+        # self.model = transformer_model.set_affine_true_for_instance_norm()
+        # self.model_name = "Norm_affine"
+        #
+        # self.model = transformer_model.replace_activation_function("GELU")
+        # self.model_name = "Activation_GELU"
+        # self.model = transformer_model.replace_activation_function("LeakyReLU")
+        # self.model_name = "Activation_LeakyRelu"
+        # self.model = transformer_model.replace_activation_function("ELU")
+        # self.model_name = "Activation_ELU"
+        # self.model = transformer_model.replace_activation_function("Mish")
+        # self.model_name = "Activation_Mish"
+        #
+        # self.model = transformer_model.add_encoder_layers(num_new_layers=2)
+        # self.model_name = "EncoderLayers_2"
+        # self.model = transformer_model.add_encoder_layers(num_new_layers=4)
+        # self.model_name = "EncoderLayers_4"
+        #
+        # self.model = transformer_model.dropout_value_change(val=0.01)
+        # self.model_name = "Dropout_01"
+        # self.model = transformer_model.dropout_value_change(val=0.001)
+        # self.model_name = "Dropout_001"
+        # self.model = transformer_model.dropout_value_change(val=0.9)
+        # self.model_name = "Dropout_9"
+        #
+        # self.model = transformer_model.change_linear_output_layers()
+        # self.model_name = "OutputLayer"
+        #
+        # self.model = transformer_model.add_cross_attention()
+        # self.model_name = "CrossAttention"
+        #
+        # self.model_name = "lr_001"
+        # self.model_name = "lr_00001"
+        #
+        # self.model_name = "wd_01"
+
+        self.model_name = ""
+
+        print(self.model)
+        self.count_parameters_by_layer()
+
+        # TODO ------------------------------------------------------------------------
+
+        self.mp_trainer = MixedPrecisionTrainer(
+            model=self.model,
+            use_fp16=self.use_fp16,
+            fp16_scale_growth=fp16_scale_growth,
+        )
+
+        self.opt = AdamW(
+            self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay
+        )
+
+        if self.resume_step:
+            self._load_optimizer_state()
+            # Model was resumed, either due to a restart or a checkpoint
+            # being specified at the command line.
+            self.ema_params = [
+                self._load_ema_parameters(rate) for rate in self.ema_rate
+            ]
+        else:
+            self.ema_params = [
+                copy.deepcopy(self.mp_trainer.master_params)
+                for _ in range(len(self.ema_rate))
+            ]
+
+        if th.cuda.is_available():
+            self.use_ddp = True
+            self.ddp_model = DDP(
+                self.model,
+                device_ids=[dist_util.dev()],
+                output_device=dist_util.dev(),
+                broadcast_buffers=False,
+                bucket_cap_mb=128,
+                find_unused_parameters=False,
+            )
+        else:
+            if dist.get_world_size() > 1:
+                logger.warn(
+                    "Distributed training requires CUDA. "
+                    "Gradients will not be synchronized properly!"
+                )
+            self.use_ddp = False
+            self.ddp_model = self.model
+
+    # TODO----------------------------------------------------------------------------------
+    def count_parameters(self):
+        model = self.model
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        untrainable_params = sum(p.numel() for p in model.parameters() if not p.requires_grad)
+
+        print(f"Trainable parameters: {trainable_params}")
+        print(f"Untrainable parameters: {untrainable_params}")
+        return trainable_params, untrainable_params
+
+    def count_parameters_by_layer(self):
+        print(f"{'Layer':<55} {'Trainable Params':<20} {'Untrainable Params':<20}")
+        print("=" * 95)
+
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                trainable_params = param.numel()
+                untrainable_params = 0
+            else:
+                trainable_params = 0
+                untrainable_params = param.numel()
+
+            print(f"{name:<55} {trainable_params:<20} {untrainable_params:<20}")
+
+        print("=" * 95)
+        total_trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total_untrainable = sum(p.numel() for p in self.model.parameters() if not p.requires_grad)
+
+        print(f"{'Total':<55} {total_trainable:<20} {total_untrainable:<20}")
+
+    def load_pretrained(self, pretrained_path):
+        state_dict = th.load(pretrained_path, map_location=dist_util.dev())
+        self.model.load_state_dict(state_dict)
+        print(self.model)
+        logger.log(f"Loaded pretrained model from {pretrained_path}")
+
+    # --------------------------------------------------------------------------------------
+
+    def _load_and_sync_parameters(self):
+        resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+
+        if resume_checkpoint:
+            self.resume_step = parse_resume_step_from_filename(resume_checkpoint)
+            # if dist.get_rank() == 0:
+            logger.log(f"loading model from checkpoint: {resume_checkpoint}...")
+            self.model.load_state_dict(
+                dist_util.load_state_dict(
+                    resume_checkpoint, map_location=dist_util.dev()
+                )
+            )
+
+        dist_util.sync_params(self.model.parameters())
+
+    def _load_ema_parameters(self, rate):
+        ema_params = copy.deepcopy(self.mp_trainer.master_params)
+
+        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+        ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate)
+        if ema_checkpoint:
+            if dist.get_rank() == 0:
+                logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...")
+                state_dict = dist_util.load_state_dict(
+                    ema_checkpoint, map_location=dist_util.dev()
+                )
+                ema_params = self.mp_trainer.state_dict_to_master_params(state_dict)
+
+        dist_util.sync_params(ema_params)
+        return ema_params
+
+    def _load_optimizer_state(self):
+        main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint
+        opt_checkpoint = bf.join(
+            bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt"
+        )
+        if bf.exists(opt_checkpoint):
+            logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}")
+            state_dict = dist_util.load_state_dict(
+                opt_checkpoint, map_location=dist_util.dev()
+            )
+            self.opt.load_state_dict(state_dict)
+
+    def run_loop(self):
+        while (
+                not self.lr_anneal_steps
+                or self.step + self.resume_step < self.lr_anneal_steps
+        ):
+            batch, cond = next(self.data)
+            self.run_step(batch, cond)
+            # TODO: change 100000 for new lr
+            if self.step % 100000 == 0:
+                lr = self.lr * (0.1 ** (self.step // 100000))
+                logger.log(f"Step {self.step}: Updating learning rate to {lr}")
+                for param_group in self.opt.param_groups:
+                    param_group["lr"] = lr
+            if self.step % self.log_interval == 0:
+                logger.dumpkvs()
+            if self.step % self.save_interval == 0 and self.step > 0:
+                self.save()
+                # Run for a finite amount of time in integration tests.
+                if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0:
+                    return
+            self.step += 1
+        # Save the last checkpoint if it wasn't already saved.
+        if (self.step - 1) % self.save_interval != 0:
+            self.save()
+
+    def run_step(self, batch, cond):
+        self.forward_backward(batch, cond)
+        took_step = self.mp_trainer.optimize(self.opt)
+        if took_step:
+            self._update_ema()
+        self._anneal_lr()
+        self.log_step()
+
+    def forward_backward(self, batch, cond):
+        self.mp_trainer.zero_grad()
+        for i in range(0, batch.shape[0], self.microbatch):
+            micro = batch[i: i + self.microbatch].to(dist_util.dev())
+            micro_cond = {
+                k: v[i: i + self.microbatch].to(dist_util.dev())
+                for k, v in cond.items()
+            }
+            model_kwargs = micro_cond
+
+            last_batch = (i + self.microbatch) >= batch.shape[0]
+            t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev())
+
+            compute_losses = functools.partial(
+                self.diffusion.training_losses,
+                self.ddp_model,
+                micro,
+                t,
+                model_kwargs=model_kwargs,
+                analog_bit=self.analog_bit,
+            )
+
+            if last_batch or not self.use_ddp:
+                losses = compute_losses()
+            else:
+                with self.ddp_model.no_sync():
+                    losses = compute_losses()
+
+            if isinstance(self.schedule_sampler, LossAwareSampler):
+                self.schedule_sampler.update_with_local_losses(
+                    t, losses["loss"].detach()
+                )
+
+            loss = (losses["loss"] * weights).mean()
+            log_loss_dict(
+                self.diffusion, t, {k: v * weights for k, v in losses.items()}
+            )
+            self.mp_trainer.backward(loss)
+
+    def _update_ema(self):
+        for rate, params in zip(self.ema_rate, self.ema_params):
+            update_ema(params, self.mp_trainer.master_params, rate=rate)
+
+    def _anneal_lr(self):
+        if not self.lr_anneal_steps:
+            return
+        frac_done = (self.step + self.resume_step) / self.lr_anneal_steps
+        lr = self.lr * (1 - frac_done)
+        for param_group in self.opt.param_groups:
+            param_group["lr"] = lr
+
+    def log_step(self):
+        logger.logkv("step", self.step + self.resume_step)
+        logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch)
+
+    def save(self):
+        def save_checkpoint(rate, params):
+            state_dict = self.mp_trainer.master_params_to_state_dict(params)
+            if dist.get_rank() == 0:
+                logger.log(f"saving model {rate}...")
+                if not rate:
+                    filename = f"model{(self.step + self.resume_step):06d}.pt"
+                else:
+                    filename = f"ema_{rate}_{(self.step + self.resume_step):06d}.pt"
+
+                filename = self.model_name + "_" + filename
+
+                with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f:
+                    th.save(state_dict, f)
+
+        save_checkpoint(0, self.mp_trainer.master_params)
+        for rate, params in zip(self.ema_rate, self.ema_params):
+            save_checkpoint(rate, params)
+
+        if dist.get_rank() == 0:
+            with bf.BlobFile(
+                    bf.join(get_blob_logdir(), f"opt{(self.step + self.resume_step):06d}.pt"),
+                    "wb",
+            ) as f:
+                th.save(self.opt.state_dict(), f)
+
+        dist.barrier()
+
+
+def parse_resume_step_from_filename(filename):
+    """
+    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
+    checkpoint's number of steps.
+    """
+    split = filename.split("model")
+    if len(split) < 2:
+        return 0
+    split1 = split[-1].split(".")[0]
+    try:
+        return int(split1)
+    except ValueError:
+        return 0
+
+
+def get_blob_logdir():
+    # You can change this to be a separate path to save checkpoints to
+    # a blobstore or some external drive.
+    return logger.get_dir()
+
+
+def find_resume_checkpoint():
+    # On your infrastructure, you may want to override this to automatically
+    # discover the latest checkpoint on your blob storage, etc.
+    return None
+
+
+def find_ema_checkpoint(main_checkpoint, step, rate):
+    if main_checkpoint is None:
+        return None
+    filename = f"ema_{rate}_{(step):06d}.pt"
+    path = bf.join(bf.dirname(main_checkpoint), filename)
+    if bf.exists(path):
+        return path
+    return None
+
+
+def log_loss_dict(diffusion, ts, losses):
+    for key, values in losses.items():
+        logger.logkv_mean(key, values.mean().item())
+        # Log the quantiles (four quartiles, in particular).
+        for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()):
+            quartile = int(4 * sub_t / diffusion.num_timesteps)
+            logger.logkv_mean(f"{key}_q{quartile}", sub_loss)

house_diffusion/transformer.py

@@ -0,0 +1,284 @@
+import math
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .nn import timestep_embedding
+
+def dec2bin(xinp, bits):
+    mask = 2 ** th.arange(bits - 1, -1, -1).to(xinp.device, xinp.dtype)
+    return xinp.unsqueeze(-1).bitwise_and(mask).ne(0).float()
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = th.arange(max_len).unsqueeze(1)
+        div_term = th.exp(th.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = th.zeros(1, max_len, d_model)
+        pe[0, :, 0::2] = th.sin(position * div_term)
+        pe[0, :, 1::2] = th.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        """
+        Args:
+            x: Tensor, shape [batch_size, seq_len, embedding_dim]
+        """
+        x = x + self.pe[0:1, :x.size(1)]
+        return self.dropout(x)
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, d_ff, dropout, activation):
+        super().__init__() 
+        # We set d_ff as a default to 2048
+        self.linear_1 = nn.Linear(d_model, d_ff)
+        self.dropout = nn.Dropout(dropout)
+        self.linear_2 = nn.Linear(d_ff, d_model)
+        self.activation = activation
+    def forward(self, x):
+        x = self.dropout(self.activation(self.linear_1(x)))
+        x = self.linear_2(x)
+        return x
+
+def attention(q, k, v, d_k, mask=None, dropout=None):
+    scores = th.matmul(q, k.transpose(-2, -1)) /  math.sqrt(d_k)
+    if mask is not None:
+        mask = mask.unsqueeze(1)
+        scores = scores.masked_fill(mask == 1, -1e9)
+    scores = F.softmax(scores, dim=-1)
+    if dropout is not None:
+        scores = dropout(scores)
+    output = th.matmul(scores, v)
+    return output
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, heads, d_model, dropout = 0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.d_k = d_model // heads
+        self.h = heads
+        self.q_linear = nn.Linear(d_model, d_model)
+        self.v_linear = nn.Linear(d_model, d_model)
+        self.k_linear = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.out = nn.Linear(d_model, d_model)
+    
+    def forward(self, q, k, v, mask=None):
+        bs = q.size(0)
+        # perform linear operation and split into h heads
+        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
+        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
+        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)
+        # transpose to get dimensions bs * h * sl * d_model
+        k = k.transpose(1,2)
+        q = q.transpose(1,2)
+        v = v.transpose(1,2)# calculate attention using function we will define next
+        #TODO
+        mask = mask.to('cuda:0')
+        scores = attention(q, k, v, self.d_k, mask, self.dropout)
+        # concatenate heads and put through final linear layer
+        concat = scores.transpose(1,2).contiguous().view(bs, -1, self.d_model)
+        output = self.out(concat)
+        return output
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, heads, dropout, activation):
+        super().__init__()
+        self.norm_1 = nn.InstanceNorm1d(d_model)
+        self.norm_2 = nn.InstanceNorm1d(d_model)
+        self.self_attn = MultiHeadAttention(heads, d_model)
+        self.door_attn = MultiHeadAttention(heads, d_model)
+        self.gen_attn = MultiHeadAttention(heads, d_model)
+        self.ff = FeedForward(d_model, d_model*2, dropout, activation)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x, door_mask, self_mask, gen_mask):
+        assert (gen_mask.max()==1 and gen_mask.min()==0), f"{gen_mask.max()}, {gen_mask.min()}"
+        x2 = self.norm_1(x)
+        x = x + self.dropout(self.door_attn(x2,x2,x2,door_mask)) \
+                + self.dropout(self.self_attn(x2, x2, x2, self_mask)) \
+                + self.dropout(self.gen_attn(x2, x2, x2, gen_mask))
+        x2 = self.norm_2(x)
+        x = x + self.dropout(self.ff(x2))
+        return x
+
+class TransformerModel(nn.Module):
+    """
+    The full Transformer model with timestep embedding.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        condition_channels,
+        model_channels,
+        out_channels,
+        dataset,
+        use_checkpoint,
+        use_unet,
+        analog_bit,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.condition_channels = condition_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.time_channels = model_channels
+        self.use_checkpoint = use_checkpoint
+        self.analog_bit = analog_bit
+        self.use_unet = use_unet
+        self.num_layers = 4
+
+        # self.pos_encoder = PositionalEncoding(model_channels, 0.001)
+        # self.activation = nn.SiLU()
+        self.activation = nn.ReLU()
+
+        self.time_embed = nn.Sequential(
+            nn.Linear(self.model_channels, self.model_channels),
+            nn.SiLU(),
+            nn.Linear(self.model_channels, self.time_channels),
+        )
+        self.input_emb = nn.Linear(self.in_channels, self.model_channels)
+        self.condition_emb = nn.Linear(self.condition_channels, self.model_channels)
+
+        if use_unet:
+            self.unet = UNet(self.model_channels, 1)
+
+        self.transformer_layers = nn.ModuleList([EncoderLayer(self.model_channels, 4, 0.1, self.activation) for x in range(self.num_layers)])
+        # self.transformer_layers = nn.ModuleList([nn.TransformerEncoderLayer(self.model_channels, 4, self.model_channels*2, 0.1, self.activation, batch_first=True) for x in range(self.num_layers)])
+
+        self.output_linear1 = nn.Linear(self.model_channels, self.model_channels)
+        self.output_linear2 = nn.Linear(self.model_channels, self.model_channels//2)
+        self.output_linear3 = nn.Linear(self.model_channels//2, self.out_channels)
+
+        if not self.analog_bit:
+            self.output_linear_bin1 = nn.Linear(162+self.model_channels, self.model_channels)
+            self.output_linear_bin2 = EncoderLayer(self.model_channels, 1, 0.1, self.activation)
+            self.output_linear_bin3 = EncoderLayer(self.model_channels, 1, 0.1, self.activation)
+            self.output_linear_bin4 = nn.Linear(self.model_channels, 16)
+
+        print(f"Number of model parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad)}")
+
+    def expand_points(self, points, connections):
+        def average_points(point1, point2):
+            points_new = (point1+point2)/2
+            return points_new
+        p1 = points
+        p1 = p1.view([p1.shape[0], p1.shape[1], 2, -1])
+        p5 = points[th.arange(points.shape[0])[:, None], connections[:,:,1].long()]
+        p5 = p5.view([p5.shape[0], p5.shape[1], 2, -1])
+        p3 = average_points(p1, p5)
+        p2 = average_points(p1, p3)
+        p4 = average_points(p3, p5)
+        p1_5 = average_points(p1, p2)
+        p2_5 = average_points(p2, p3)
+        p3_5 = average_points(p3, p4)
+        p4_5 = average_points(p4, p5)
+        points_new = th.cat((p1.view_as(points), p1_5.view_as(points), p2.view_as(points),
+            p2_5.view_as(points), p3.view_as(points), p3_5.view_as(points), p4.view_as(points), p4_5.view_as(points), p5.view_as(points)), 2)
+        return points_new.detach()
+
+    def create_image(self, points, connections, room_indices, img_size=256, res=200):
+        img = th.zeros((points.shape[0], 1, img_size, img_size), device=points.device)
+        points = (points+1)*(img_size//2)
+        points[points>=img_size] = img_size-1
+        points[points<0] = 0
+        p1 = points
+        p2 = points[th.arange(points.shape[0])[:, None], connections[:,:,1].long()]
+
+        slope = (p2[:,:,1]-p1[:,:,1])/((p2[:,:,0]-p1[:,:,0]))
+        slope[slope.isnan()] = 0
+        slope[slope.isinf()] = 1
+
+        m = th.linspace(0, 1, res, device=points.device)
+        new_shape = [p2.shape[0], res, p2.shape[1], p2.shape[2]]
+
+        new_p2 = p2.unsqueeze(1).expand(new_shape)
+        new_p1 = p1.unsqueeze(1).expand(new_shape)
+        new_room_indices = room_indices.unsqueeze(1).expand([p2.shape[0], res, p2.shape[1], 1])
+
+        inc = new_p2 - new_p1
+
+        xs =  m.view(1,-1,1) * inc[:,:,:,0]
+        xs = xs + new_p1[:,:,:,0]
+        xs = xs.long()
+
+        x_inc = th.where(inc[:,:,:,0]==0, inc[:,:,:,1], inc[:,:,:,0])
+        x_inc =  m.view(1,-1,1) * x_inc
+        ys = x_inc * slope.unsqueeze(1) + new_p1[:,:,:,1]
+        ys = ys.long()
+
+        img[th.arange(xs.shape[0])[:, None], :, xs.view(img.shape[0], -1), ys.view(img.shape[0], -1)] = new_room_indices.reshape(img.shape[0], -1, 1).float()
+        return img.detach()
+
+    def forward(self, x, timesteps, xtalpha, epsalpha, is_syn=False, **kwargs):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x S x C] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :return: an [N x S x C] Tensor of outputs.
+        """
+        # prefix = 'syn_' if is_syn else '' 
+        prefix = 'syn_' if is_syn else '' 
+        x = x.permute([0, 2, 1]).float() # -> convert [N x C x S] to [N x S x C]
+
+        if not self.analog_bit:
+            x = self.expand_points(x, kwargs[f'{prefix}connections'])
+
+        # Different input embeddings (Input, Time, Conditions)
+        #TODO---------------------------------------------------------------
+        x = x.to('cuda:0')
+        timesteps = timesteps.to(x.device)
+        # print(x.device)
+
+        time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+        time_emb = time_emb.unsqueeze(1)
+        input_emb = self.input_emb(x)
+        if self.condition_channels>0:
+            cond = None
+            for key in [f'{prefix}room_types', f'{prefix}corner_indices', f'{prefix}room_indices']:
+                if cond is None:
+                    cond = kwargs[key]
+                else:
+                    cond = th.cat((cond, kwargs[key]), 2)
+            #TODO
+            cond = cond.to('cuda:0')
+            cond_emb = self.condition_emb(cond.float())
+
+        # PositionalEncoding and DM model
+        out = input_emb + cond_emb + time_emb.repeat((1, input_emb.shape[1], 1))
+        for layer in self.transformer_layers:
+            out = layer(out, kwargs[f'{prefix}door_mask'], kwargs[f'{prefix}self_mask'], kwargs[f'{prefix}gen_mask'])
+
+        out_dec = self.output_linear1(out)
+        out_dec = self.activation(out_dec)
+        out_dec = self.output_linear2(out_dec)
+        out_dec = self.output_linear3(out_dec)
+
+        if not self.analog_bit:
+            out_bin_start = x*xtalpha.repeat([1,1,9]) - out_dec.repeat([1,1,9]) * epsalpha.repeat([1,1,9])
+            out_bin = (out_bin_start/2 + 0.5) # -> [0,1]
+            out_bin = out_bin * 256 #-> [0, 256]
+            out_bin = dec2bin(out_bin.round().int(), 8)
+            out_bin_inp = out_bin.reshape([x.shape[0], x.shape[1], 16*9])
+            out_bin_inp[out_bin_inp==0] = -1
+
+            out_bin = th.cat((out_bin_start, out_bin_inp, cond_emb), 2)
+            out_bin = self.activation(self.output_linear_bin1(out_bin))
+            out_bin = self.output_linear_bin2(out_bin, kwargs[f'{prefix}door_mask'], kwargs[f'{prefix}self_mask'], kwargs[f'{prefix}gen_mask'])
+            out_bin = self.output_linear_bin3(out_bin, kwargs[f'{prefix}door_mask'], kwargs[f'{prefix}self_mask'], kwargs[f'{prefix}gen_mask'])
+            out_bin = self.output_linear_bin4(out_bin)
+
+            out_bin = out_bin.permute([0, 2, 1]) # -> convert back [N x S x C] to [N x C x S]
+
+        out_dec = out_dec.permute([0, 2, 1]) # -> convert back [N x S x C] to [N x C x S]
+
+        if not self.analog_bit:
+            return out_dec, out_bin
+        else:
+            return out_dec, None

house_diffusion/transformer_models.py

@@ -0,0 +1,228 @@
+import torch.nn as nn
+from .transformer import TransformerModel, EncoderLayer
+
+
+# class TransformerModels(nn.Module):
+class TransformerModels:
+    def __init__(self, model, device):
+        self.model = model
+        self.device = device
+
+    """ ------------------------------- 1) Normalize ------------------------------- """
+
+    def replace_InstanceNorm1d_LayerNorm(self):
+        self.freeze_unfreeze(True)
+        for name, layer in self.model.named_modules():
+            if isinstance(layer, nn.InstanceNorm1d):
+                num_features = layer.num_features
+                new_layer = nn.LayerNorm(normalized_shape=num_features).to(self.device)
+                parent_module = dict(self.model.named_modules())[name.rsplit('.', 1)[0]]
+                setattr(parent_module, name.split('.')[-1], new_layer)
+
+        return self.model
+
+    def set_affine_true_for_instance_norm(self):
+        self.freeze_unfreeze(True)
+        for name, layer in self.model.named_modules():
+            if isinstance(layer, nn.InstanceNorm1d):
+                new_layer = nn.InstanceNorm1d(num_features=100, affine=True).to(self.device)
+                parent_module = dict(self.model.named_modules())[name.rsplit('.', 1)[0]]
+                setattr(parent_module, name.split('.')[-1], new_layer)
+
+        return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ -------------------------- 2) Activation Function -------------------------- """
+
+    def replace_activation_function(self, activation):
+        self.freeze_unfreeze(True)
+        functions = {
+            "GELU": nn.GELU(),
+            "LeakyReLU": nn.LeakyReLU(),
+            "ELU": nn.ELU(),
+            "Mish": nn.Mish(),
+            # "ReLU": nn.ReLU(),
+        }
+
+        def replace_activation_in_module(module, activation_layer):
+            for name, child in module.named_children():
+                if isinstance(child, nn.ReLU):
+                    setattr(module, name, activation_layer)
+                else:
+                    replace_activation_in_module(child, activation_layer)
+
+        new_activation_layer = functions[activation].to(self.device)
+        replace_activation_in_module(self.model, new_activation_layer)
+        return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ ---------------------------- 3) New Encoder Layers ------------------------- """
+
+    def add_encoder_layers(self, num_new_layers=2):
+        self.freeze_unfreeze(True)
+        new_encoder_layers = [EncoderLayer(512, 4, 0.1, nn.ReLU()).to(self.device) for _ in range(num_new_layers)]
+
+        for i, new_layer in enumerate(new_encoder_layers):
+            self.model.transformer_layers.insert(4 + i, new_layer.to(self.device))
+
+        return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ -------------------------------- 4) Dropout -------------------------------- """
+
+    # def dropout_value_change(self, val=0.1):
+    #     self.freeze_unfreeze(True)
+    #     for layer in self.model.modules():
+    #         if isinstance(layer, nn.Dropout):
+    #             layer.p = val
+    #
+    #     return self.model
+
+    def dropout_value_change(self, val=0.1):
+        self.freeze_unfreeze(True)
+
+        def replace_dropouts_in_module(module, rate):
+            for name, child in module.named_children():
+                if isinstance(child, nn.Dropout):
+                    setattr(module, name, nn.Dropout(rate).to(self.device))
+                else:
+                    replace_dropouts_in_module(child, rate)
+
+        replace_dropouts_in_module(self.model, val)
+
+        return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ ------------------------- 5) Output linear layers -------------------------- """
+
+    def change_linear_output_layers(self):
+        output_layers_names = [
+            "output_linear1",
+            "output_linear2",
+            "output_linear3",
+            "output_linear_bin1",
+            "output_linear_bin2",
+            "output_linear_bin3",
+        ]
+        for name, param in self.model.named_parameters():
+            param.requires_grad = False
+            if name.split(".")[0] in output_layers_names:
+                param.requires_grad = True
+
+        output_linear1 = self.model.output_linear1
+        output_linear2 = self.model.output_linear2
+        output_linear3 = self.model.output_linear3
+        output_linear_bin1 = self.model.output_linear_bin1
+        output_linear_bin2 = self.model.output_linear_bin2
+        output_linear_bin3 = self.model.output_linear_bin3
+
+        output_linear11 = nn.Linear(output_linear1.out_features,
+                                    output_linear1.out_features).to(self.device)
+        output_linear21 = nn.Linear(output_linear2.out_features,
+                                    output_linear2.out_features).to(self.device)
+
+        # self.model.output_layers = nn.Sequential(
+        #     output_linear1,
+        #     output_linear11,
+        #     output_linear2,
+        #     output_linear21,
+        #     output_linear3,
+        #     output_linear_bin1,
+        #     output_linear_bin2,
+        #     output_linear_bin3,
+        # )
+        self.model.insert(6, output_linear11)
+        self.model.insert(8, output_linear21)
+
+        return self.model
+
+    # def change_linear_output_layers(self):
+    #     output_layers_names = [
+    #         "output_linear1",
+    #         "output_linear2",
+    #         "output_linear3",
+    #         "output_linear_bin1",
+    #         "output_linear_bin2",
+    #         "output_linear_bin3",
+    #     ]
+    #     for name, param in self.model.named_parameters():
+    #         param.requires_grad = False
+    #         if name.split(".")[0] in output_layers_names:
+    #             param.requires_grad = True
+    #
+    #     output_linear1 = self.model.output_linear1
+    #     output_linear2 = self.model.output_linear2
+    #     output_linear3 = self.model.output_linear3
+    #     # output_linear_bin1 = self.model.output_linear_bin1
+    #     # output_linear_bin2 = self.model.output_linear_bin2
+    #     # output_linear_bin3 = self.model.output_linear_bin3
+    #
+    #     output_linear11 = nn.Linear(output_linear1.out_features,
+    #                                 output_linear1.out_features).to(self.device)
+    #     output_linear21 = nn.Linear(output_linear2.out_features,
+    #                                 output_linear2.out_features).to(self.device)
+    #
+    #     self.model.output_linear1.append(output_linear11.to(self.device))
+    #     self.model.output_linear2.append(output_linear21.to(self.device))
+    #
+    #     # self.model.output_layers = nn.Sequential(
+    #     #     output_linear1,
+    #     #     output_linear11,
+    #     #     output_linear2,
+    #     #     output_linear21,
+    #     #     output_linear3,
+    #     #     output_linear_bin1,
+    #     #     output_linear_bin2,
+    #     #     output_linear_bin3,
+    #     # )
+    #
+    #     return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ ---------------------------- 6) Cross-Attention ---------------------------- """
+
+    def add_cross_attention(self, embed_dim=512, num_heads=8, dropout=0.1):
+        self.freeze_unfreeze(True)
+        for idx, layer in enumerate(self.model.transformer_layers):
+            cross_attn_layer = CrossAttentionLayer(embed_dim, num_heads, dropout).to(self.device)
+            layer.gen_attn = nn.Sequential(layer.gen_attn, cross_attn_layer).to(self.device)
+
+        return self.model
+
+    """ ---------------------------------------------------------------------------- """
+    """ -------------------------- 7) Residual Connections? ------------------------- """
+
+    """ ---------------------------------------------------------------------------- """
+    """ ------------------------------- 8) Attention Heads? (check if works with same params) ------------------------------- """
+
+    """ ---------------------------------------------------------------------------- """
+    #Add LayerNorm Before/After Attention
+
+    # ADAM ?
+    # weight decay ?
+    # learning rate?
+
+    def freeze_unfreeze(self, flag):
+        for param in self.model.parameters():
+            param.requires_grad = flag
+
+    def count_parameters(self):
+        model = self.model
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        untrainable_params = sum(p.numel() for p in model.parameters() if not p.requires_grad)
+
+        print(f"Trainable parameters: {trainable_params}")
+        print(f"Untrainable parameters: {untrainable_params}")
+        return trainable_params, untrainable_params
+
+
+class CrossAttentionLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.1):
+        super(CrossAttentionLayer, self).__init__()
+        self.cross_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, query, key_value, attn_mask=None):
+        attn_output, _ = self.cross_attn(query, key_value, key_value, attn_mask=attn_mask)
+        return self.norm(self.dropout(attn_output) + query)