import streamlit as st
from streamlit_drawable_canvas import st_canvas
from huggingface_hub import HfFileSystem
from flax.serialization import msgpack_restore, from_state_dict

fs = HfFileSystem()

with fs.open("PrakhAI/HelloWorld/checkpoint.msgpack", "rb") as f:
    state = msgpack_restore(f.read())

print(type(state))
print(state)
print([(k, type(v)) for (k, v) in state.items()])
# print(state['params'])
# print(state['opt_state'])

# logits = state.apply_fn({'params': state.params}, batch['image'])
# print(logits)

# x = st.slider('Select a value')
# st.write(x, 'squared is', x * x)

# Specify brush parameters and drawing mode
b_width = st.sidebar.slider("Brush width: ", 1, 100, 10)
b_color = st.beta_color_picker("Enter brush color hex: ")
bg_color = st.beta_color_picker("Enter background color hex: ", "#eee")
drawing_mode = st.sidebar.checkbox("Drawing mode ?", True)

# Create a canvas component
image_data = st_canvas(
    b_width, b_color, bg_color, height=150, drawing_mode=drawing_mode, key="canvas"
)