Update app.py

Added dit_size specific num_samples limit

app.py

@@ -166,11 +166,27 @@ def generate_random_seed():
     """Generate a random seed between 0 and 2^32 - 1"""
     return random.randint(0, 2**32 - 1)
 
+MODEL_SAMPLE_LIMITS = {
+    "S": {"min":1, "max": 18, "default": 4},
+    "B": {"min":1, "max": 9, "default": 4},
+    "L": {"min":1, "max": 3, "default": 1}
+}
+
+def update_sample_slider(dit_size):
+    limits = MODEL_SAMPLE_LIMITS[dit_size]
+    return gr.update(
+        minimum=limits["min"],
+        maximum=limits["max"],
+        value=limits["default"],
+        info=f"How many images to generate ({limits['min']}-{limits['max']})"        
+    )
+
 @spaces.GPU(duration=120)
 def generate_samples(num_samples, dit_size, genre_name, style_name, seed, progress=gr.Progress()):
     """Main function for Gradio interface"""
-    if num_samples < 1 or num_samples > 12:
-        return None, gr.update(value="Number of samples must be between 1 and 12", visible=True)
+    limits = MODEL_SAMPLE_LIMITS[dit_size]
+    if num_samples < limits["min"] or num_samples > limits["max"]:
+        return None, gr.update(value=f"Number of samples for {dit_size} model must be between {limits["min"]} and {limits["max"]}", visible=True)
     
     # Get genre and style IDs from mappings
     genre_id = reduced_genre_mapping.get(genre_name)
@@ -205,8 +221,21 @@ with gr.Blocks(title="DiT Diffusion Model Generator", theme=gr.themes.Soft()) as
     
     with gr.Row():
         with gr.Column(scale=1):
-            num_samples = gr.Slider(minimum=1, maximum=12, value=2, step=1, label="Number of Samples", info="How many images to generate (1-12)")
-            dit_size = gr.Radio(choices=["S", "B", "L"], value="B", label="DiT Model Size", info="Larger models produce better quality but take longer. Recommend to choose <4 num_samples for L(Large) model.")
+            dit_size = gr.Radio(
+                choices=["S", "B", "L"], 
+                value="B", 
+                label="DiT Model Size", 
+                info="S: Small (fastest), B: Base (balanced), L: Large (best quality but slowest)"
+            )
+
+            num_samples = gr.Slider(
+                minimum=MODEL_SAMPLE_LIMITS["B"]["min"],
+                maximum=MODEL_SAMPLE_LIMITS["B"]["max"],
+                value=MODEL_SAMPLE_LIMITS["B"]["default"],
+                step=1,
+                label="Number of Samples",
+                info=f"How many images to generate ({MODEL_SAMPLE_LIMITS['B']['min']}-{MODEL_SAMPLE_LIMITS['B']['max']})"
+            )
             
             genre_names = list(reduced_genre_mapping.keys())
             style_names = list(reduced_style_mapping.keys())
@@ -230,8 +259,10 @@ with gr.Blocks(title="DiT Diffusion Model Generator", theme=gr.themes.Soft()) as
             progress_bar = gr.Progress(track_tqdm=True)
             
         with gr.Column(scale=2):
-            output_gallery = gr.Gallery(label="Generated Images", columns=4, rows=3, object_fit="contain", height=600)
+            output_gallery = gr.Gallery(label="Generated Images", columns=4, rows=4, object_fit="contain", height=600)
             error_message = gr.Textbox(label="Error", visible=False, max_lines=3, container=True, elem_id="error_box")
+
+    dit_size.change(update_sample_slider, inputs=[dit_size],outputs=[num_samples])
     
     # Seed reset button functionality
     reset_seed_btn.click(generate_random_seed, inputs=[], outputs=[seed])