import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
from SparseCAE import SparseCAE

model = SparseCAE.load_from_checkpoint("best.ckpt")
model.eval()

def decode(z0, z1, z2, z3):
    z = torch.tensor([[z0, z1, z2, z3]], dtype=torch.float32)
    with torch.no_grad():
        img = model.decoder(z).numpy().reshape(28, 28)
    fig, ax = plt.subplots(figsize=(3, 3))
    ax.imshow(img, cmap="gray"); ax.axis("off")
    return fig

gr.Interface(
    fn=decode,
    inputs=[gr.Slider(0, 1, step=0.01, label=f"z[{i}]") for i in range(4)],
    outputs=gr.Plot(),
    live=True,
).launch()