Update app.py

app.py

@@ -72,7 +72,7 @@ class DiffusionSampler:
             self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema").to(self.device)
             self.vae.eval()
 
-    @spaces.GPU
+    @spaces.GPU(duration=120)
     def generate_images(self, model, num_samples, genre, style, seed, progress=gr.Progress()):
         """Generate images with the DiT model"""
         global global_progress
@@ -166,11 +166,11 @@ def generate_random_seed():
     """Generate a random seed between 0 and 2^32 - 1"""
     return random.randint(0, 2**32 - 1)
 
-@spaces.GPU
+@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 > 16:
-        return None, gr.update(value="Number of samples must be between 1 and 16", visible=True)
+    if num_samples < 1 or num_samples > 12:
+        return None, gr.update(value="Number of samples must be between 1 and 12", visible=True)
     
     # Get genre and style IDs from mappings
     genre_id = reduced_genre_mapping.get(genre_name)
@@ -205,8 +205,8 @@ 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=16, value=4, step=1, label="Number of Samples", info="How many images to generate (1-16)")
-            dit_size = gr.Radio(choices=["S", "B", "L"], value="S", label="DiT Model Size", info="Larger models produce better quality but take longer")
+            num_samples = gr.Slider(minimum=1, maximum=12, value=2, step=1, label="Number of Samples", info="How many images to generate (1-16)")
+            dit_size = gr.Radio(choices=["S", "B", "L"], value="B", label="DiT Model Size", info="Larger models produce better quality but take longer")
             
             genre_names = list(reduced_genre_mapping.keys())
             style_names = list(reduced_style_mapping.keys())
@@ -230,7 +230,7 @@ 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=4, object_fit="contain", height=600)
+            output_gallery = gr.Gallery(label="Generated Images", columns=4, rows=3, object_fit="contain", height=600)
             error_message = gr.Textbox(label="Error", visible=False, max_lines=3, container=True, elem_id="error_box")
     
     # Seed reset button functionality