add adjustable size

app.py

@@ -5,7 +5,7 @@ import zipfile
 import shutil
 from test_api import run_adain
 
-def process_adain_api(files, style_strength=1.0):
+def process_adain_api(files, style_strength=1.0, dataset_size=100):
     '''API endpoint for AdaIN processing.'''
     if not files:
         return None, "No files uploaded"
@@ -27,7 +27,7 @@ def process_adain_api(files, style_strength=1.0):
         style_dataset_pth = "tidalove/paleo-real"
         
         # Run AdaIN
-        result_path = run_adain(input_dir, style_dataset_pth, output_dir, style_strength)
+        run_adain(input_dir, style_dataset_pth, output_dir, style_strength, dataset_size)
         
         # Create zip file
         zip_path = os.path.join(temp_dir, "style_transfer_results.zip")
@@ -48,7 +48,8 @@ with gr.Blocks() as adain_demo:
     with gr.Row():
         with gr.Column():
             files_input = gr.File(label="Upload Images", file_count="multiple", file_types=["image"])
-            strength_input = gr.Slider(0.0, 2.0, 1.0, step=0.1, label="Style Strength")
+            strength_input = gr.Slider(0.0, 2.0, 1.0, step=0.1, label="Style strength")
+            size_input = gr.Slider(1, 1000, 100, step=10, label="Generated dataset size")
             process_btn = gr.Button("Process", variant="primary")
             
         with gr.Column():
@@ -57,7 +58,7 @@ with gr.Blocks() as adain_demo:
     
     process_btn.click(
         fn=process_adain_api,
-        inputs=[files_input, strength_input],
+        inputs=[files_input, strength_input, size_input],
         outputs=[download_output, status_output],
         api_name="adain_process"  # This creates the API endpoint
     )