p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics_diffusion.py (python)

    gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
    metrics = dict(
        psnr=psnr,
        ssim=ssim,
    )
[content] jasmine/train_tokenizer_vqvae.py (python)
            q_loss = jnp.square(
                jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]
            ).mean()
            commitment_loss = jnp.square(
                outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
            ).mean()
            loss = mse + q_loss + args.vq_beta * commitment_loss
            _, index_counts = jnp.unique_counts(
                jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
            )
            codebook_usage = (index_counts != 0).mean()
            metrics["loss"] = loss
            metrics["q_loss"] = q_loss