Update app.py

app.py

@@ -12,7 +12,6 @@ print(f'Using device: {device}')
 # Create a scheduler
 noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')
 
-
 class ClassConditionedUnet(nn.Module):
   def __init__(self, num_classes=10, class_emb_size=4):
     super().__init__()
@@ -90,7 +89,7 @@ def load_checkpoint_for_inference(filepath, model_class):
 
 
 
-def generate_images(selected_class_name, num_samples=4):
+def generate_images(selected_class_name, num_samples):
     print(f"Generating {num_samples} samples for class: {selected_class_name}")
 
     # Map class name to class ID
@@ -130,16 +129,15 @@ def generate_images(selected_class_name, num_samples=4):
 
 
 
-
-# Initialize a dummy model (replace with your actual model loading)
-model = load_checkpoint_for_inference(filepath="CIFAR10_unet_v_02_100_epochs_inference.pth", model_class=ClassConditionedUnet)
-
 # CIFAR-10 class names
 cifar10_classes = [
     "plane", "car", "bird", "cat", "deer",
     "dog", "frog", "horse", "ship", "truck"
 ]
 
+# Initialize a dummy model (replace with your actual model loading)
+model = load_checkpoint_for_inference(filepath="CIFAR10_unet_v_02_100_epochs_inference.pth", model_class=ClassConditionedUnet)
+
 
 # Create the Gradio interface
 custom_css = """
@@ -173,16 +171,19 @@ with gr.Blocks(css=custom_css) as demo:
         class_selector = gr.Radio(
             cifar10_classes, label="Select CIFAR-10 Class", value=cifar10_classes[0]
         )
+        num_samples_slider = gr.Slider(
+            minimum=1, maximum=10, value=4, step=1, label="Number of Images to Generate"
+        )
 
     with gr.Row():
-        generate_btn = gr.Button("Generate Samples")
+        generate_btn = gr.Button("Generate")
 
     with gr.Row():
         output_gallery = gr.Gallery(label="Generated Images", show_label=True, elem_id="gallery")
 
     generate_btn.click(
         fn=generate_images,
-        inputs=class_selector,
+        inputs=[class_selector, num_samples_slider],
         outputs=output_gallery
     )