Update app.py

app.py

@@ -5,7 +5,8 @@ import tensorflow as tf
 import gradio as gr
 import numpy as np
 
-model = from_pretrained_keras("keras-io/WGAN-GP")
+model1 = tf.keras.models.load_model("mnist.h5", compile=False)
+model2 = from_pretrained_keras("keras-io/WGAN-GP")
 
 title = "WGAN-GP"
 description = "Image Generation Using WGAN"
@@ -17,7 +18,7 @@ article = """
         </p>
         """
 
-def Predict(num_images):
+def Predict(model, num_images):
         random_latent_vectors = tf.random.normal(shape=(int(num_images), 128))
         preds = model(random_latent_vectors) 
         num = ceil(sqrt(num_images))
@@ -32,13 +33,21 @@ def Predict(num_images):
         return images
 
 
+def inference(num_images, select: str):
+    if select == 'fmnist':
+        result = create_digit_samples(model2, num_images)
+    else:
+        result = create_digit_samples(model1, num_images)
+    return result
+    
 examples = [[5],[8],[2],[3]]
-
+inputs = [gr.inputs.Number(label="number of images"), gr.inputs.Radio(['fmnist', 'mnist'])]
+outputs = gr.outputs.Image(label="Output Image")
 
 interface = gr.Interface(
-    fn = Predict,
-    inputs = ["number"],
-    outputs = ["image"],
+    fn = inference,
+    inputs = inputs,
+    outputs = outputs,
     examples = examples,
     description = description,
     title = title,