Gholamreza's picture
Upload 5 files
8895b4f verified
raw
history blame contribute delete
974 Bytes
import torch
import gradio as gr
from models import Generator
from conditional_gan import generate_digit
generator = Generator()
def init():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the generator
generator.load_state_dict(torch.load('models/generator.pt', map_location=device))
generator.to(device)
def generate_mnist_digit(digit):
return generate_digit(generator, digit)
# Gradio Interface
def gradio_generate(digit):
return generate_mnist_digit(digit)
with gr.Blocks() as demo:
gr.Markdown("# MNIST Digit Generator")
digit = gr.Dropdown(list(range(10)), label="Select a Digit")
generate_button = gr.Button("Generate")
output_image = gr.Image(label="Generated Image", type="filepath")
generate_button.click(gradio_generate, inputs=digit, outputs=output_image)
if __name__ == '__main__':
init()
print("* Model loaded")
demo.launch()