Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
from diffusers.models import UNet2DModel | |
from huggingface_hub import hf_hub_download | |
path = hf_hub_download(repo_id="porestar/oadg_channels_64", filename="model.pt") | |
model = UNet2DModel( | |
sample_size=64, | |
in_channels=2, | |
out_channels=2, | |
layers_per_block=2, | |
block_out_channels=(64, 64, 128, 128), | |
down_block_types=( | |
"DownBlock2D", | |
"DownBlock2D", | |
"AttnDownBlock2D", | |
"DownBlock2D", | |
), | |
up_block_types=( | |
"UpBlock2D", | |
"AttnUpBlock2D", | |
"UpBlock2D", | |
"UpBlock2D", | |
), | |
) | |
model.load_state_dict(torch.load(path)) | |
def classify_image(inp): | |
return {"lol": 0} | |
img = gr.Image(image_mode="L", source="canvas", shape=(32, 32), invert_colors=False) | |
label = gr.Label(num_top_classes=3) | |
demo = gr.Interface( | |
fn=classify_image, inputs=img, outputs=label, interpretation="default" | |
) | |
demo.launch() | |