import gradio as gr
import lightning as L

from src.gradio_demo.color_matching import create_color_matching
from src.gradio_demo.sw_guidance import create_sw_guidance

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # ReSWD

        **[Project Page](https://reservoirswd.github.io/)** | 
        **[Paper](https://arxiv.org/abs/2510.01061)** |
        **[GitHub](https://github.com/Stability-AI/ReSWD)**

        ReSTIR‘d, not shaken. Combining Reservoir Sampling and Sliced Wasserstein
        Distance for Variance Reduction.

        ReSWD is a method for distribution matching with reduced variance.
        """
    )
    with gr.Tab("SW Guidance (SD 3.5 Large Turbo)"):
        create_sw_guidance("stabilityai/stable-diffusion-3.5-large-turbo")
    with gr.Tab("Color Matching"):
        create_color_matching()

    demo.launch()