Make a download zone

We can't creazte a zip file and start the downloading...
So, let's create a download zone with gr.File for the moment.

Also, cleanup and type annotation.

app.py

@@ -8,29 +8,35 @@ from packager import create_dataset_zip
 from documentation import DOC_CONTENT
 
 
-def auto_label(image, imid, dataset):
+def auto_label(
+    image: Image.Image, imid: int, dataset: ImageDataset
+) -> tuple[str, ImageDataset]:
     text = get_task_response("<MORE_DETAILED_CAPTION>", image)
     ds = dataset.update_label(imid, text)
     return text, ds
 
 
-def label_changed(label, imid, dataset):
+def label_changed(label: str, imid: int, dataset: ImageDataset) -> ImageDataset:
     return dataset.update_label(imid, label)
 
 
-def update_single_label(dataset, label_text, image_id):
+def update_single_label(
+    dataset: ImageDataset, label_text: str, image_id: int
+) -> ImageDataset:
     """Update single image label in dataset."""
     return dataset.update_label(image_id, label_text)
 
 
-def uploaded(files, current_dataset):
+def uploaded(files: list, current_dataset: ImageDataset | None) -> ImageDataset:
     """Handle file upload - return new dataset instance."""
     if current_dataset is None:
         current_dataset = ImageDataset()
     return current_dataset.add_images(files)
 
 
-def labelize_all_images(dataset, label, progress=gr.Progress(True)):
+def labelize_all_images(
+    dataset: ImageDataset, label: str, progress=gr.Progress(True)
+) -> tuple[ImageDataset, str]:
     """Generate labels for all images and return new dataset instance."""
 
     # Generate actual labels
@@ -42,13 +48,13 @@ def labelize_all_images(dataset, label, progress=gr.Progress(True)):
     return dataset.update_all_labels(labels_dict), label
 
 
-def download_dataset(dataset, organize_in_folders):
+def create_dataset_zipfile(dataset: ImageDataset, organize_in_folders: bool):
     """Create and return zip file for download."""
     zip_path = create_dataset_zip(dataset, organize_in_folders)
-    return zip_path
+    return gr.update(visible=True, value=zip_path)
 
 
-def update_buttons_states(dataset, labeling_in_progress=False):
+def update_buttons_states(dataset: ImageDataset, labeling_in_progress=False):
     """Update all button states based on dataset and labeling progress."""
     count = len(dataset.images)
     return (
@@ -61,12 +67,12 @@ def update_buttons_states(dataset, labeling_in_progress=False):
     )
 
 
-def start_labeling(dataset):
+def start_labeling(dataset: ImageDataset):
     """Start labeling process - disable buttons and show progress."""
     return update_buttons_states(dataset, labeling_in_progress=True)
 
 
-def finish_labeling(dataset):
+def finish_labeling(dataset: ImageDataset):
     """Finish labeling process - enable buttons and hide progress."""
     return update_buttons_states(dataset, labeling_in_progress=False)
 
@@ -78,17 +84,38 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
         with gr.Group():
             gr.Markdown("Upload images and add labels to build your dataset.")
 
-            upload_button = gr.UploadButton("📁 Upload images", file_count="multiple")
-            label_all = gr.Button("🏷️ Labelize all images", interactive=False)
-            is_labeling_in_progress = gr.State(False)
-            progressbar = gr.Label("", visible=False, label="Preparing...")
-            remove_all = gr.Button("🗑️ Remove all", interactive=False)
+            upload_button = gr.UploadButton(
+                "📁 Upload images",
+                file_count="multiple",
+            )
+            label_all = gr.Button(
+                "🏷️ Labelize all images",
+                interactive=False,
+            )
+            is_labeling_in_progress = gr.State(
+                False,
+            )
+            progressbar = gr.Label(
+                "",
+                visible=False,
+                label="Preparing...",
+            )
+            remove_all = gr.Button(
+                "🗑️ Remove all",
+                interactive=False,
+            )
 
         with gr.Group():
-            download_button = gr.DownloadButton(
-                "💾 Download dataset", interactive=False, size="lg"
+            organize_files = gr.Checkbox(
+                label="📂 Organize in folders", value=False, render=False
+            )
+            download_button = gr.Button(
+                "💾 Create zip file to download",
+                interactive=False,
+                size="lg",
             )
-            organize_files = gr.Checkbox(label="📂 Organize in folders", value=False)
+            download_file = gr.File(label="Generated datasets", visible=False)
+            organize_files.render()
 
     @gr.render(inputs=[dataset, is_labeling_in_progress])
     def render_grid(ds, is_labeling_in_progress):
@@ -97,7 +124,6 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
             return
 
         # Hidden component to trigger label refresh
-        gr.Number(value=0, visible=False)
         with gr.Row(equal_height=True):
             for im in ds.images:
                 with (
@@ -112,6 +138,9 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
                         ],
                     ),
                 ):
+                    # Hidden component to store current image ID
+                    current_image_id = gr.State(value=im["id"])
+
                     image = gr.Image(
                         im["path"],
                         type="pil",
@@ -123,6 +152,7 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
                     )
 
                     label = gr.Text(
+                        im["label"],
                         placeholder="Description...",
                         lines=5,
                         container=False,
@@ -130,9 +160,6 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
                         key=f"text_{im['id']}",
                     )
 
-                    # Hidden component to store current image ID
-                    current_image_id = gr.State(value=im["id"])
-
                     button = gr.Button(
                         "✨ Generate label",
                         interactive=not is_labeling_in_progress,
@@ -218,12 +245,15 @@ with gr.Blocks(title="Labelizer", fill_width=True) as demo:
             is_labeling_in_progress,
         ],
     )
-
-    # Download zip file
+    # create the zip file and set the download file section ready to use
     download_button.click(
-        download_dataset,
+        lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=download_file,
+    ).then(
+        create_dataset_zipfile,
         inputs=[dataset, organize_files],
-        outputs=[download_button],
+        outputs=download_file,
     )