add feedback collection functionality through HF OAuth

README.md

@@ -10,6 +10,8 @@ pinned: true
 license: cc-by-sa-4.0
 models:
 - prs-eth/marigold-lcm-v1-0
+hf_oauth: true
+hf_oauth_expiration_minutes: 43200
 ---
 
 This is a demo of Marigold-LCM, the state-of-the-art depth estimator for images in the wild.

app.py

@@ -16,18 +16,19 @@
 # Please find bibtex at: https://github.com/prs-eth/Marigold#-citation
 # More information about the method can be found at https://marigoldmonodepth.github.io
 # --------------------------------------------------------------------------
-
+from __future__ import annotations
 
 import functools
 import os
 import tempfile
+import warnings
 import zipfile
 from io import BytesIO
 
-import spaces
 import gradio as gr
 import imageio as imageio
 import numpy as np
+import spaces
 import torch as torch
 from PIL import Image
 from gradio_imageslider import ImageSlider
@@ -35,8 +36,13 @@ from huggingface_hub import login
 from tqdm import tqdm
 
 from extrude import extrude_depth_3d
+from flagging import FlagMethod, HuggingFaceDatasetSaver
 from marigold_depth_estimation_lcm import MarigoldDepthConsistencyPipeline
 
+warnings.filterwarnings(
+    "ignore", message=".*LoginButton created outside of a Blocks context.*"
+)
+
 default_seed = 2024
 
 default_image_denoise_steps = 4
@@ -64,6 +70,16 @@ default_bas_frame_thickness = 5
 default_bas_frame_near = 1
 default_bas_frame_far = 1
 
+default_share_always_show_hf_logout_btn = True
+default_share_always_show_accordion = False
+
+
+def process_image_check(path_input):
+    if path_input is None:
+        raise gr.Error(
+            "Missing image in the first pane: upload a file or use one from the gallery below."
+        )
+
 
 def process_image(
     pipe,
@@ -72,6 +88,14 @@ def process_image(
     ensemble_size=default_image_ensemble_size,
     processing_res=default_image_processing_res,
 ):
+    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing image {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
+    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy")
+    path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png")
+    path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png")
+
     input_image = Image.open(path_input)
 
     pipe_out = pipe(
@@ -88,13 +112,6 @@ def process_image(
     depth_colored = pipe_out.depth_colored
     depth_16bit = (depth_pred * 65535.0).astype(np.uint16)
 
-    path_output_dir = tempfile.mkdtemp()
-
-    name_base = os.path.splitext(os.path.basename(path_input))[0]
-    path_out_fp32 = os.path.join(path_output_dir, f"{name_base}_depth_fp32.npy")
-    path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.png")
-    path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.png")
-
     np.save(path_out_fp32, depth_pred)
     Image.fromarray(depth_16bit).save(path_out_16bit, mode="I;16")
     depth_colored.save(path_out_vis)
@@ -116,9 +133,15 @@ def process_video(
     out_max_frames=default_video_out_max_frames,
     progress=gr.Progress(),
 ):
-    path_output_dir = tempfile.mkdtemp()
+    if path_input is None:
+        raise gr.Error(
+            "Missing video in the first pane: upload a file or use one from the gallery below."
+        )
 
-    name_base = os.path.splitext(os.path.basename(path_input))[0]
+    name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing video {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
     path_out_vis = os.path.join(path_output_dir, f"{name_base}_depth_colored.mp4")
     path_out_16bit = os.path.join(path_output_dir, f"{name_base}_depth_16bit.zip")
 
@@ -218,12 +241,18 @@ def process_bas(
     frame_near=default_bas_frame_near,
     frame_far=default_bas_frame_far,
 ):
+    if path_input is None:
+        raise gr.Error(
+            "Missing image in the first pane: upload a file or use one from the gallery below."
+        )
+
     if plane_near >= plane_far:
         raise gr.Error("NEAR plane must have a value smaller than the FAR plane")
 
-    path_output_dir = tempfile.mkdtemp()
-
     name_base, name_ext = os.path.splitext(os.path.basename(path_input))
+    print(f"Processing bas-relief {name_base}{name_ext}")
+
+    path_output_dir = tempfile.mkdtemp()
 
     input_image = Image.open(path_input)
 
@@ -267,9 +296,11 @@ def process_bas(
         path_glb, path_stl = extrude_depth_3d(
             image_rgb_new,
             image_depth_new,
-            output_model_scale=size_longest_cm * 10
-            if output_model_scale is None
-            else output_model_scale,
+            output_model_scale=(
+                size_longest_cm * 10
+                if output_model_scale is None
+                else output_model_scale
+            ),
             filter_size=filter_size,
             coef_near=plane_near,
             coef_far=plane_far,
@@ -288,17 +319,22 @@ def process_bas(
         256, filter_size, vertex_colors=False, scene_lights=True, output_model_scale=1
     )
     path_files_glb, path_files_stl = _process_3d(
-        size_longest_px, filter_size, vertex_colors=True, scene_lights=False, prepare_for_3d_printing=True
+        size_longest_px,
+        filter_size,
+        vertex_colors=True,
+        scene_lights=False,
+        prepare_for_3d_printing=True,
     )
 
     return path_viewer_glb, [path_files_glb, path_files_stl]
 
 
-def run_demo_server(pipe):
+def run_demo_server(pipe, hf_writer=None):
     process_pipe_image = spaces.GPU(functools.partial(process_image, pipe))
-    process_pipe_video = spaces.GPU(functools.partial(process_video, pipe), duration=120)
+    process_pipe_video = spaces.GPU(
+        functools.partial(process_video, pipe), duration=120
+    )
     process_pipe_bas = spaces.GPU(functools.partial(process_bas, pipe))
-    os.environ["GRADIO_ALLOW_FLAGGING"] = "never"
 
     gradio_theme = gr.themes.Default()
 
@@ -332,6 +368,9 @@ def run_demo_server(pipe):
                 text-align: center;
                 display: block;
             }
+            .md_feedback li {
+                margin-bottom: 0px !important;
+            }
         """,
         head="""
             <script async src="https://www.googletagmanager.com/gtag/js?id=G-1FWSVCGZTG"></script>
@@ -343,35 +382,70 @@ def run_demo_server(pipe):
             </script>
         """,
     ) as demo:
+        if hf_writer is not None:
+            print("Creating login button")
+            share_login_btn = gr.LoginButton(size="sm", scale=1, render=False)
+            print("Created login button")
+            share_login_btn.activate()
+            print("Activated login button")
+
         gr.Markdown(
             """
             # Marigold-LCM Depth Estimation
             <p align="center">
-            <a title="Website" href="https://marigoldmonodepth.github.io/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <a title="Website" href="https://marigoldmonodepth.github.io/" target="_blank" rel="noopener noreferrer" 
+                    style="display: inline-block;">
                 <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
             </a>
-            <a title="arXiv" href="https://arxiv.org/abs/2312.02145" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <a title="arXiv" href="https://arxiv.org/abs/2312.02145" target="_blank" rel="noopener noreferrer" 
+                    style="display: inline-block;">
                 <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
             </a>
-            <a title="Github" href="https://github.com/prs-eth/marigold" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                <img src="https://img.shields.io/github/stars/prs-eth/marigold?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+            <a title="Github" href="https://github.com/prs-eth/marigold" target="_blank" rel="noopener noreferrer" 
+                    style="display: inline-block;">
+                <img src="https://img.shields.io/github/stars/prs-eth/marigold?label=GitHub%20%E2%98%85&logo=github&color=C8C" 
+                        alt="badge-github-stars">
             </a>
-            <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+            <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" 
+                    style="display: inline-block;">
                 <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
             </a>
             </p>
             <p align="justify">
-                Marigold-LCM is the fast version of Marigold, the state-of-the-art depth estimator for images in the wild.
-                It combines the power of the original Marigold 10-step estimator and the Latent Consistency Models, delivering high-quality results in as little as <b>one step</b>.
-                We provide three functions in this demo: Image, Video, and Bas-relief 3D processing — <b>see the tabs below</b>. 
-                Upload your content into the <b>first</b> pane, or click any of the <b>examples</b> below.
-                Wait a second (for images and 3D) or a minute (for videos), and interact with the result in the <b>second</b> pane.
-                To avoid queuing, fork the demo into your profile.
-                <a href="https://huggingface.co/spaces/prs-eth/marigold">The original Marigold demo is also available</a>.
+                Marigold-LCM is the fast version of Marigold, the state-of-the-art depth estimator for images in the 
+                wild. It combines the power of the original Marigold 10-step estimator and the Latent Consistency 
+                Models, delivering high-quality results in as little as <b>one step</b>. We provide three functions 
+                in this demo: Image, Video, and Bas-relief 3D processing — <b>see the tabs below</b>. Upload your 
+                content into the <b>first</b> pane, or click any of the <b>examples</b> below. Wait a second (for 
+                images and 3D) or a minute (for videos), and interact with the result in the <b>second</b> pane. To 
+                avoid queuing, fork the demo into your profile.
+                <a href="https://huggingface.co/spaces/prs-eth/marigold">
+                    The original Marigold demo is also available
+                </a>.
             </p>
         """
         )
 
+        def get_share_instructions(is_full):
+            out = (
+                "### Help us improve Marigold! If the output is not what you expected, "
+                "you can help us by sharing it with us privately.\n"
+            )
+            if is_full:
+                out += (
+                    "1. Sign into your Hugging Face account using the button below.\n"
+                    "1. Signing in may reset the demo and results; in that case, process the image again.\n"
+                )
+            out += "1. Review and agree to the terms of usage and enter an optional message to us.\n"
+            out += "1. Click the 'Share' button to submit the image to us privately.\n"
+            return out
+
+        def get_share_conditioned_on_login(profile: gr.OAuthProfile | None):
+            state_logged_out = profile is None
+            return get_share_instructions(is_full=state_logged_out), gr.Button(
+                visible=(state_logged_out or default_share_always_show_hf_logout_btn)
+            )
+
         with gr.Tabs(elem_classes=["tabs"]):
             with gr.Tab("Image"):
                 with gr.Row():
@@ -423,6 +497,42 @@ def run_demo_server(pipe):
                             elem_id="download",
                             interactive=False,
                         )
+
+                        if hf_writer is not None:
+                            with gr.Accordion(
+                                "Feedback",
+                                open=False,
+                                visible=default_share_always_show_accordion,
+                            ) as share_box:
+                                share_instructions = gr.Markdown(
+                                    get_share_instructions(is_full=True),
+                                    elem_classes="md_feedback",
+                                )
+                                share_transfer_of_rights = gr.Checkbox(
+                                    label="(Optional) I own or hold necessary rights to the submitted image. By "
+                                    "checking this box, I grant an irrevocable, non-exclusive, transferable, "
+                                    "royalty-free, worldwide license to use the uploaded image, including for "
+                                    "publishing, reproducing, and model training. [transfer_of_rights]",
+                                    scale=1,
+                                )
+                                share_content_is_legal = gr.Checkbox(
+                                    label="By checking this box, I acknowledge that my uploaded content is legal and "
+                                    "safe, and that I am solely responsible for ensuring it complies with all "
+                                    "applicable laws and regulations. Additionally, I am aware that my Hugging Face "
+                                    "username is collected. [content_is_legal]",
+                                    scale=1,
+                                )
+                                share_reason = gr.Textbox(
+                                    label="(Optional) Reason for feedback",
+                                    max_lines=1,
+                                    interactive=True,
+                                )
+                                with gr.Row():
+                                    share_login_btn.render()
+                                    share_share_btn = gr.Button(
+                                        "Share", variant="stop", scale=1
+                                    )
+
                 gr.Examples(
                     fn=process_pipe_image,
                     examples=[
@@ -502,7 +612,8 @@ def run_demo_server(pipe):
                     """
                     <p align="justify">
                         This part of the demo uses Marigold-LCM to create a bas-relief model. 
-                        The models are watertight, with correct normals, and exported in the STL format, which makes them <b>3D-printable</b>.
+                        The models are watertight, with correct normals, and exported in the STL format, which makes 
+                        them <b>3D-printable</b>.
                     </p>
                     """,
                 )
@@ -513,7 +624,9 @@ def run_demo_server(pipe):
                             type="filepath",
                         )
                         with gr.Row():
-                            bas_submit_btn = gr.Button(value="Create 3D", variant="primary")
+                            bas_submit_btn = gr.Button(
+                                value="Create 3D", variant="primary"
+                            )
                             bas_reset_btn = gr.Button(value="Reset")
                         with gr.Accordion("3D printing demo: Main options", open=True):
                             bas_plane_near = gr.Slider(
@@ -537,7 +650,9 @@ def run_demo_server(pipe):
                                 step=1,
                                 value=default_bas_embossing,
                             )
-                        with gr.Accordion("3D printing demo: Advanced options", open=False):
+                        with gr.Accordion(
+                            "3D printing demo: Advanced options", open=False
+                        ):
                             bas_denoise_steps = gr.Slider(
                                 label="Number of denoising steps",
                                 minimum=1,
@@ -682,17 +797,66 @@ def run_demo_server(pipe):
                     cache_examples=True,
                 )
 
-        image_submit_btn.click(
-            fn=process_pipe_image,
-            inputs=[
-                image_input,
-                image_denoise_steps,
-                image_ensemble_size,
-                image_processing_res,
-            ],
-            outputs=[image_output_slider, image_output_files],
-            concurrency_limit=1,
-        )
+        ### Image tab
+
+        if hf_writer is not None:
+            image_submit_btn.click(
+                fn=process_image_check,
+                inputs=image_input,
+                outputs=None,
+                preprocess=False,
+                queue=False,
+            ).success(
+                get_share_conditioned_on_login,
+                None,
+                [share_instructions, share_login_btn],
+                queue=False,
+            ).then(
+                lambda: (
+                    gr.Button(value="Share", interactive=True),
+                    gr.Accordion(visible=True),
+                    False,
+                    False,
+                    "",
+                ),
+                None,
+                [
+                    share_share_btn,
+                    share_box,
+                    share_transfer_of_rights,
+                    share_content_is_legal,
+                    share_reason,
+                ],
+                queue=False,
+            ).then(
+                fn=process_pipe_image,
+                inputs=[
+                    image_input,
+                    image_denoise_steps,
+                    image_ensemble_size,
+                    image_processing_res,
+                ],
+                outputs=[image_output_slider, image_output_files],
+                concurrency_limit=1,
+            )
+        else:
+            image_submit_btn.click(
+                fn=process_image_check,
+                inputs=image_input,
+                outputs=None,
+                preprocess=False,
+                queue=False,
+            ).success(
+                fn=process_pipe_image,
+                inputs=[
+                    image_input,
+                    image_denoise_steps,
+                    image_ensemble_size,
+                    image_processing_res,
+                ],
+                outputs=[image_output_slider, image_output_files],
+                concurrency_limit=1,
+            )
 
         image_reset_btn.click(
             fn=lambda: (
@@ -712,9 +876,73 @@ def run_demo_server(pipe):
                 image_denoise_steps,
                 image_processing_res,
             ],
-            concurrency_limit=1,
+            queue=False,
         )
 
+        if hf_writer is not None:
+            image_reset_btn.click(
+                fn=lambda: (
+                    gr.Button(value="Share", interactive=True),
+                    gr.Accordion(visible=default_share_always_show_accordion),
+                ),
+                inputs=[],
+                outputs=[
+                    share_share_btn,
+                    share_box,
+                ],
+                queue=False,
+            )
+
+        ### Share functionality
+
+        if hf_writer is not None:
+            share_components = [
+                image_input,
+                image_denoise_steps,
+                image_ensemble_size,
+                image_processing_res,
+                image_output_slider,
+                share_content_is_legal,
+                share_transfer_of_rights,
+                share_reason,
+            ]
+
+            hf_writer.setup(share_components, "shared_data")
+            share_callback = FlagMethod(hf_writer, "Share", "", visual_feedback=True)
+
+            def share_precheck(
+                hf_content_is_legal,
+                image_output_slider,
+                profile: gr.OAuthProfile | None,
+            ):
+                if profile is None:
+                    raise gr.Error(
+                        "Log into the Space with your Hugging Face account first."
+                    )
+                if image_output_slider is None or image_output_slider[0] is None:
+                    raise gr.Error("No output detected; process the image first.")
+                if not hf_content_is_legal:
+                    raise gr.Error(
+                        "You must consent that the uploaded content is legal."
+                    )
+                return gr.Button(value="Sharing in progress", interactive=False)
+
+            share_share_btn.click(
+                share_precheck,
+                [share_content_is_legal, image_output_slider],
+                share_share_btn,
+                preprocess=False,
+                queue=False,
+            ).success(
+                share_callback,
+                inputs=share_components,
+                outputs=share_share_btn,
+                preprocess=False,
+                queue=False,
+            )
+
+        ### Video tab
+
         video_submit_btn.click(
             fn=process_pipe_video,
             inputs=[video_input],
@@ -729,6 +957,8 @@ def run_demo_server(pipe):
             concurrency_limit=1,
         )
 
+        ### Bas-relief tab
+
         bas_submit_btn.click(
             fn=process_pipe_bas,
             inputs=[
@@ -791,6 +1021,8 @@ def run_demo_server(pipe):
             concurrency_limit=1,
         )
 
+        ### Server launch
+
         demo.queue(
             api_open=False,
         ).launch(
@@ -801,6 +1033,7 @@ def run_demo_server(pipe):
 
 def main():
     CHECKPOINT = "prs-eth/marigold-lcm-v1-0"
+    CROWD_DATA = "crowddata-marigold-lcm-v1-0-space-v1-0"
 
     if "HF_TOKEN_LOGIN" in os.environ:
         login(token=os.environ["HF_TOKEN_LOGIN"])
@@ -816,7 +1049,18 @@ def main():
         pass  # run without xformers
 
     pipe = pipe.to(device)
-    run_demo_server(pipe)
+
+    hf_writer = None
+    if "HF_TOKEN_LOGIN" in os.environ:
+        hf_writer = HuggingFaceDatasetSaver(
+            os.getenv("HF_TOKEN_LOGIN"),
+            CROWD_DATA,
+            private=True,
+            info_filename="dataset_info.json",
+            separate_dirs=True,
+        )
+
+    run_demo_server(pipe, hf_writer)
 
 
 if __name__ == "__main__":

extrude.py

@@ -336,7 +336,9 @@ def extrude_depth_3d(
     mesh.apply_scale(scaling_factor)
 
     if prepare_for_3d_printing:
-        rotation_mat = trimesh.transformations.rotation_matrix(np.radians(90), [-1, 0, 0])
+        rotation_mat = trimesh.transformations.rotation_matrix(
+            np.radians(90), [-1, 0, 0]
+        )
         mesh.apply_transform(rotation_mat)
 
     path_out_base = os.path.splitext(path_depth)[0].replace("_16bit", "")

flagging.py

@@ -0,0 +1,387 @@
+from __future__ import annotations
+
+import csv
+import json
+import time
+import uuid
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import filelock
+import huggingface_hub
+from gradio_client import utils as client_utils
+from gradio_client.documentation import document
+
+import gradio as gr
+from gradio import utils
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+
+
+class FlaggingCallback(ABC):
+    """
+    An abstract class for defining the methods that any FlaggingCallback should have.
+    """
+
+    @abstractmethod
+    def setup(self, components: list[Component], flagging_dir: str):
+        """
+        This method should be overridden and ensure that everything is set up correctly for flag().
+        This method gets called once at the beginning of the Interface.launch() method.
+        Parameters:
+        components: Set of components that will provide flagged data.
+        flagging_dir: A string, typically containing the path to the directory where the flagging file should be stored (provided as an argument to Interface.__init__()).
+        """
+        pass
+
+    @abstractmethod
+    def flag(
+        self,
+        flag_data: list[Any],
+        flag_option: str = "",
+        username: str | None = None,
+    ) -> int:
+        """
+        This method should be overridden by the FlaggingCallback subclass and may contain optional additional arguments.
+        This gets called every time the <flag> button is pressed.
+        Parameters:
+        interface: The Interface object that is being used to launch the flagging interface.
+        flag_data: The data to be flagged.
+        flag_option (optional): In the case that flagging_options are provided, the flag option that is being used.
+        username (optional): The username of the user that is flagging the data, if logged in.
+        Returns:
+        (int) The total number of samples that have been flagged.
+        """
+        pass
+
+
+@document()
+class HuggingFaceDatasetSaver(FlaggingCallback):
+    """
+    A callback that saves each flagged sample (both the input and output data) to a HuggingFace dataset.
+
+    Example:
+        import gradio as gr
+        hf_writer = gr.HuggingFaceDatasetSaver(HF_API_TOKEN, "image-classification-mistakes")
+        def image_classifier(inp):
+            return {'cat': 0.3, 'dog': 0.7}
+        demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label",
+                            allow_flagging="manual", flagging_callback=hf_writer)
+    Guides: using-flagging
+    """
+
+    def __init__(
+        self,
+        hf_token: str,
+        dataset_name: str,
+        private: bool = False,
+        info_filename: str = "dataset_info.json",
+        separate_dirs: bool = False,
+    ):
+        """
+        Parameters:
+            hf_token: The HuggingFace token to use to create (and write the flagged sample to) the HuggingFace dataset (defaults to the registered one).
+            dataset_name: The repo_id of the dataset to save the data to, e.g. "image-classifier-1" or "username/image-classifier-1".
+            private: Whether the dataset should be private (defaults to False).
+            info_filename: The name of the file to save the dataset info (defaults to "dataset_infos.json").
+            separate_dirs: If True, each flagged item will be saved in a separate directory. This makes the flagging more robust to concurrent editing, but may be less convenient to use.
+        """
+        self.hf_token = hf_token
+        self.dataset_id = dataset_name  # TODO: rename parameter (but ensure backward compatibility somehow)
+        self.dataset_private = private
+        self.info_filename = info_filename
+        self.separate_dirs = separate_dirs
+
+    def setup(self, components: list[Component], flagging_dir: str):
+        """
+        Params:
+        flagging_dir (str): local directory where the dataset is cloned,
+        updated, and pushed from.
+        """
+        # Setup dataset on the Hub
+        self.dataset_id = huggingface_hub.create_repo(
+            repo_id=self.dataset_id,
+            token=self.hf_token,
+            private=self.dataset_private,
+            repo_type="dataset",
+            exist_ok=True,
+        ).repo_id
+        path_glob = "**/*.jsonl" if self.separate_dirs else "data.csv"
+        huggingface_hub.metadata_update(
+            repo_id=self.dataset_id,
+            repo_type="dataset",
+            metadata={
+                "configs": [
+                    {
+                        "config_name": "default",
+                        "data_files": [{"split": "train", "path": path_glob}],
+                    }
+                ]
+            },
+            overwrite=True,
+            token=self.hf_token,
+        )
+
+        # Setup flagging dir
+        self.components = components
+        self.dataset_dir = (
+            Path(flagging_dir).absolute() / self.dataset_id.split("/")[-1]
+        )
+        self.dataset_dir.mkdir(parents=True, exist_ok=True)
+        self.infos_file = self.dataset_dir / self.info_filename
+
+        # Download remote files to local
+        remote_files = [self.info_filename]
+        if not self.separate_dirs:
+            # No separate dirs => means all data is in the same CSV file => download it to get its current content
+            remote_files.append("data.csv")
+
+        for filename in remote_files:
+            try:
+                huggingface_hub.hf_hub_download(
+                    repo_id=self.dataset_id,
+                    repo_type="dataset",
+                    filename=filename,
+                    local_dir=self.dataset_dir,
+                    token=self.hf_token,
+                )
+            except huggingface_hub.utils.EntryNotFoundError:
+                pass
+
+    def flag(
+        self,
+        flag_data: list[Any],
+        flag_option: str = "",
+        username: str | None = None,
+    ) -> int:
+        if self.separate_dirs:
+            # JSONL files to support dataset preview on the Hub
+            current_utc_time = datetime.now(timezone.utc)
+            iso_format_without_microseconds = current_utc_time.strftime(
+                "%Y-%m-%dT%H:%M:%S"
+            )
+            milliseconds = int(current_utc_time.microsecond / 1000)
+            unique_id = f"{iso_format_without_microseconds}.{milliseconds:03}Z"
+            if username not in (None, ""):
+                unique_id += f"_U_{username}"
+            else:
+                unique_id += f"_{str(uuid.uuid4())[:8]}"
+            components_dir = self.dataset_dir / unique_id
+            data_file = components_dir / "metadata.jsonl"
+            path_in_repo = unique_id  # upload in sub folder (safer for concurrency)
+        else:
+            # Unique CSV file
+            components_dir = self.dataset_dir
+            data_file = components_dir / "data.csv"
+            path_in_repo = None  # upload at root level
+
+        return self._flag_in_dir(
+            data_file=data_file,
+            components_dir=components_dir,
+            path_in_repo=path_in_repo,
+            flag_data=flag_data,
+            flag_option=flag_option,
+            username=username or "",
+        )
+
+    def _flag_in_dir(
+        self,
+        data_file: Path,
+        components_dir: Path,
+        path_in_repo: str | None,
+        flag_data: list[Any],
+        flag_option: str = "",
+        username: str = "",
+    ) -> int:
+        # Deserialize components (write images/audio to files)
+        features, row = self._deserialize_components(
+            components_dir, flag_data, flag_option, username
+        )
+
+        # Write generic info to dataset_infos.json + upload
+        with filelock.FileLock(str(self.infos_file) + ".lock"):
+            if not self.infos_file.exists():
+                self.infos_file.write_text(
+                    json.dumps({"flagged": {"features": features}})
+                )
+
+                huggingface_hub.upload_file(
+                    repo_id=self.dataset_id,
+                    repo_type="dataset",
+                    token=self.hf_token,
+                    path_in_repo=self.infos_file.name,
+                    path_or_fileobj=self.infos_file,
+                )
+
+        headers = list(features.keys())
+
+        if not self.separate_dirs:
+            with filelock.FileLock(components_dir / ".lock"):
+                sample_nb = self._save_as_csv(data_file, headers=headers, row=row)
+                sample_name = str(sample_nb)
+                huggingface_hub.upload_folder(
+                    repo_id=self.dataset_id,
+                    repo_type="dataset",
+                    commit_message=f"Flagged sample #{sample_name}",
+                    path_in_repo=path_in_repo,
+                    ignore_patterns="*.lock",
+                    folder_path=components_dir,
+                    token=self.hf_token,
+                )
+        else:
+            sample_name = self._save_as_jsonl(data_file, headers=headers, row=row)
+            sample_nb = len(
+                [path for path in self.dataset_dir.iterdir() if path.is_dir()]
+            )
+            huggingface_hub.upload_folder(
+                repo_id=self.dataset_id,
+                repo_type="dataset",
+                commit_message=f"Flagged sample #{sample_name}",
+                path_in_repo=path_in_repo,
+                ignore_patterns="*.lock",
+                folder_path=components_dir,
+                token=self.hf_token,
+            )
+
+        return sample_nb
+
+    @staticmethod
+    def _save_as_csv(data_file: Path, headers: list[str], row: list[Any]) -> int:
+        """Save data as CSV and return the sample name (row number)."""
+        is_new = not data_file.exists()
+
+        with data_file.open("a", newline="", encoding="utf-8") as csvfile:
+            writer = csv.writer(csvfile)
+
+            # Write CSV headers if new file
+            if is_new:
+                writer.writerow(utils.sanitize_list_for_csv(headers))
+
+            # Write CSV row for flagged sample
+            writer.writerow(utils.sanitize_list_for_csv(row))
+
+        with data_file.open(encoding="utf-8") as csvfile:
+            return sum(1 for _ in csv.reader(csvfile)) - 1
+
+    @staticmethod
+    def _save_as_jsonl(data_file: Path, headers: list[str], row: list[Any]) -> str:
+        """Save data as JSONL and return the sample name (uuid)."""
+        Path.mkdir(data_file.parent, parents=True, exist_ok=True)
+        with open(data_file, "w") as f:
+            json.dump(dict(zip(headers, row)), f)
+        return data_file.parent.name
+
+    def _deserialize_components(
+        self,
+        data_dir: Path,
+        flag_data: list[Any],
+        flag_option: str = "",
+        username: str = "",
+    ) -> tuple[dict[Any, Any], list[Any]]:
+        """Deserialize components and return the corresponding row for the flagged sample.
+
+        Images/audio are saved to disk as individual files.
+        """
+        # Components that can have a preview on dataset repos
+        file_preview_types = {gr.Audio: "Audio", gr.Image: "Image"}
+
+        # Generate the row corresponding to the flagged sample
+        features = OrderedDict()
+        row = []
+        for component, sample in zip(self.components, flag_data):
+            # Get deserialized object (will save sample to disk if applicable -file, audio, image,...-)
+            label = component.label or ""
+            save_dir = data_dir / client_utils.strip_invalid_filename_characters(label)
+            save_dir.mkdir(exist_ok=True, parents=True)
+            deserialized = component.flag(sample, save_dir)
+
+            # Base component .flag method returns JSON; extract path from it when it is FileData
+            if component.data_model:
+                data = component.data_model.from_json(json.loads(deserialized))
+                if component.data_model == gr.data_classes.FileData:
+                    deserialized = data.path
+
+            # Add deserialized object to row
+            features[label] = {"dtype": "string", "_type": "Value"}
+            try:
+                deserialized_path = Path(deserialized)
+                if not deserialized_path.exists():
+                    raise FileNotFoundError(f"File {deserialized} not found")
+                row.append(str(deserialized_path.relative_to(self.dataset_dir)))
+            except (FileNotFoundError, TypeError, ValueError):
+                deserialized = "" if deserialized is None else str(deserialized)
+                row.append(deserialized)
+
+            # If component is eligible for a preview, add the URL of the file
+            # Be mindful that images and audio can be None
+            if isinstance(component, tuple(file_preview_types)):  # type: ignore
+                for _component, _type in file_preview_types.items():
+                    if isinstance(component, _component):
+                        features[label + " file"] = {"_type": _type}
+                        break
+                if deserialized:
+                    path_in_repo = str(  # returned filepath is absolute, we want it relative to compute URL
+                        Path(deserialized).relative_to(self.dataset_dir)
+                    ).replace(
+                        "\\", "/"
+                    )
+                    row.append(
+                        huggingface_hub.hf_hub_url(
+                            repo_id=self.dataset_id,
+                            filename=path_in_repo,
+                            repo_type="dataset",
+                        )
+                    )
+                else:
+                    row.append("")
+        features["flag"] = {"dtype": "string", "_type": "Value"}
+        features["username"] = {"dtype": "string", "_type": "Value"}
+        row.append(flag_option)
+        row.append(username)
+        return features, row
+
+
+class FlagMethod:
+    """
+    Helper class that contains the flagging options and calls the flagging method. Also
+    provides visual feedback to the user when flag is clicked.
+    """
+
+    def __init__(
+        self,
+        flagging_callback: FlaggingCallback,
+        label: str,
+        value: str,
+        visual_feedback: bool = True,
+    ):
+        self.flagging_callback = flagging_callback
+        self.label = label
+        self.value = value
+        self.__name__ = "Flag"
+        self.visual_feedback = visual_feedback
+
+    def __call__(
+        self,
+        request: gr.Request,
+        profile: gr.OAuthProfile | None,
+        *flag_data,
+    ):
+        username = None
+        if profile is not None:
+            username = profile.username
+        try:
+            self.flagging_callback.flag(
+                list(flag_data), flag_option=self.value, username=username
+            )
+        except Exception as e:
+            print(f"Error while sharing: {e}")
+            if self.visual_feedback:
+                return gr.Button(value="Sharing error", interactive=False)
+        if not self.visual_feedback:
+            return
+        time.sleep(0.8)  # to provide enough time for the user to observe button change
+        return gr.Button(value="Sharing complete", interactive=False)

marigold_depth_estimation_lcm.py

@@ -391,7 +391,9 @@ class MarigoldDepthConsistencyPipeline(DiffusionPipeline):
             ).sample  # [B, 4, h, w]
 
             # compute the previous noisy sample x_t -> x_t-1
-            depth_latent = self.scheduler.step(noise_pred, t, depth_latent, generator=rng).prev_sample
+            depth_latent = self.scheduler.step(
+                noise_pred, t, depth_latent, generator=rng
+            ).prev_sample
 
         depth = self._decode_depth(depth_latent)
 

requirements.txt

@@ -1,16 +1,16 @@
 gradio==4.21.0
-gradio-imageslider==0.0.16
+gradio-imageslider==0.0.18
 pygltflib==1.16.1
 trimesh==4.0.5
 imageio
 imageio-ffmpeg
 Pillow
 
-spaces>=0.25.0
-accelerate>=0.22.0
+spaces==0.25.0
+accelerate==0.25.0
 diffusers==0.27.2
 matplotlib==3.8.2
 scipy==1.11.4
 torch==2.0.1
-transformers>=4.32.1
-xformers>=0.0.21
+transformers==4.36.1
+xformers==0.0.21