initial

.gitignore

@@ -0,0 +1,211 @@
+outputs
+.gradio
+
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

app.py

@@ -0,0 +1,294 @@
+import os
+import spaces
+import gradio as gr
+import torch
+import functools
+import numpy as np
+import torch.nn.functional as F
+from diffusers import FluxPipeline, AutoencoderTiny
+from transformers import CLIPProcessor, CLIPModel, AutoModel
+from transformers.models.clip.modeling_clip import _get_vector_norm
+from my_utils.group_inference import run_group_inference
+from my_utils.default_values import apply_defaults
+import argparse
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1").to("cuda")
+
+m_clip = CLIPModel.from_pretrained("multimodalart/clip-vit-base-patch32").to("cuda")
+prep_clip = CLIPProcessor.from_pretrained("multimodalart/clip-vit-base-patch32")
+dino_model = AutoModel.from_pretrained('facebook/dinov2-base').to("cuda")
+
+# Get default args for flux-schnell
+default_args = argparse.Namespace(
+    model_name="flux-schnell",
+    prompt=None,
+    starting_candidates=None,
+    output_group_size=None,
+    pruning_ratio=None,
+    lambda_score=None,
+    seed=None,
+    unary_term="clip_text_img",
+    binary_term="diversity_dino",
+    guidance_scale=None,
+    num_inference_steps=None,
+    height=None,
+    width=None,
+)
+default_args = apply_defaults(default_args)
+
+
+# Scoring functions
+@torch.no_grad()
+def unary_clip_text_img_score(l_images, target_caption, device="cuda"):
+    """Compute CLIP text-image similarity scores."""
+    _img_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(device)
+    _img_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(device)
+    
+    b_images = torch.cat(l_images, dim=0)
+    b_images = F.interpolate(b_images, size=(224, 224), mode="bilinear", align_corners=False)
+    b_images = b_images * 0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    
+    text_encoding = prep_clip.tokenizer(target_caption, return_tensors="pt", padding=True).to(device)
+    output = m_clip(pixel_values=b_images, **text_encoding).logits_per_image / m_clip.logit_scale.exp()
+    return output.view(-1).cpu().numpy()
+
+
+@torch.no_grad()
+def binary_dino_diversity_score(l_images, device="cuda"):
+    """Compute pairwise diversity scores using DINO."""
+    b_images = torch.cat(l_images, dim=0)
+    _img_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    _img_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    
+    b_images = F.interpolate(b_images, size=(256, 256), mode="bilinear", align_corners=False)
+    b_images = b_images * 0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    all_features = dino_model(pixel_values=b_images).last_hidden_state[:, 1:, :].cpu()
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = all_features[i]
+        for j in range(i+1, N):
+            f2 = all_features[j]
+            cos_sim = (1 - F.cosine_similarity(f1, f2, dim=1)).mean().item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix
+
+
+@torch.no_grad()
+def binary_dino_cls_score(l_images, device="cuda"):
+    """Compute pairwise diversity scores using DINO CLS tokens."""
+    b_images = torch.cat(l_images, dim=0)
+    _img_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    _img_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    
+    b_images = F.interpolate(b_images, size=(256, 256), mode="bilinear", align_corners=False)
+    b_images = b_images * 0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    all_features = dino_model(pixel_values=b_images).last_hidden_state[:, 0:1, :].cpu()
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = all_features[i]
+        for j in range(i+1, N):
+            f2 = all_features[j]
+            cos_sim = (1 - F.cosine_similarity(f1, f2, dim=1)).mean().item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix
+
+
+@torch.no_grad()
+def binary_clip_diversity_score(l_images, device="cuda"):
+    """Compute pairwise diversity scores using CLIP."""
+    _img_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(device)
+    _img_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(device)
+    
+    b_images = torch.cat(l_images, dim=0)
+    b_images = F.interpolate(b_images, size=(224, 224), mode="bilinear", align_corners=False)
+    b_images = b_images * 0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    
+    vision_outputs = m_clip.vision_model(
+        pixel_values=b_images, 
+        output_attentions=False, 
+        output_hidden_states=False,
+        interpolate_pos_encoding=False, 
+        return_dict=True
+    )
+    image_embeds = m_clip.visual_projection(vision_outputs[1])
+    image_embeds = image_embeds / _get_vector_norm(image_embeds)
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = image_embeds[i]
+        for j in range(i+1, N):
+            f2 = image_embeds[j]
+            cos_sim = (1 - torch.dot(f1, f2)).item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix
+
+
+def get_score_functions(unary_term, binary_term, prompt):
+    """Get the appropriate scoring functions based on selected terms."""
+    # Unary score function (always CLIP for flux-schnell) - bind the prompt
+    unary_score_fn = functools.partial(unary_clip_text_img_score, target_caption=prompt, device="cuda")
+    # Binary score function
+    if binary_term == "diversity_dino":
+        binary_score_fn = functools.partial(binary_dino_diversity_score, device="cuda")
+    elif binary_term == "dino_cls_pairwise":
+        binary_score_fn = functools.partial(binary_dino_cls_score, device="cuda")
+    elif binary_term == "diversity_clip":
+        binary_score_fn = functools.partial(binary_clip_diversity_score, device="cuda")
+    else:
+        raise ValueError(f"Invalid binary term: {binary_term}")
+    
+    return unary_score_fn, binary_score_fn
+
+
+@spaces.GPU(duration=300)
+def generate_images(prompt, starting_candidates, output_group_size, pruning_ratio, 
+                   lambda_score, seed, unary_term, binary_term, progress=gr.Progress(track_tqdm=True)):
+    """Generate images using group inference with progressive pruning."""
+    
+    # Get scoring functions with prompt bound to unary function
+    unary_score_fn, binary_score_fn = get_score_functions(unary_term, binary_term, prompt)
+    
+    # Create inference args
+    inference_args = {
+        "model_name": "flux-schnell",
+        "prompt": prompt,
+        "guidance_scale": default_args.guidance_scale,
+        "num_inference_steps": default_args.num_inference_steps,
+        "max_sequence_length": 256,
+        "height": default_args.height,
+        "width": default_args.width,
+        "unary_score_fn": unary_score_fn,
+        "binary_score_fn": binary_score_fn,
+        "output_group_size": output_group_size,
+        "pruning_ratio": pruning_ratio,
+        "lambda_score": lambda_score,
+        "l_generator": [torch.Generator("cpu").manual_seed(seed + i) for i in range(starting_candidates)],
+        "starting_candidates": starting_candidates,
+        "skip_first_cfg": True,
+    }
+    print(f"pruning ratio is: {pruning_ratio}")
+    # Run group inference
+    output_group = run_group_inference(pipe, **inference_args)
+    return output_group
+
+
+# Load custom CSS
+css_path = os.path.join(os.path.dirname(__file__), "styles.css")
+with open(css_path, "r") as f:
+    custom_css = f.read()
+
+# JavaScript to force light mode
+js_func = """
+function refresh() {
+    const url = new URL(window.location);
+    if (url.searchParams.get('__theme') !== 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.href = url.href;
+    }
+}
+"""
+
+# Create Gradio interface
+with gr.Blocks(css=custom_css, js=js_func, theme=gr.themes.Soft(), elem_id="main-container") as demo:
+    
+    # Title and header
+    gr.HTML(
+        """
+        <div class="title_left">
+        <h1>Scaling Group Inference for Diverse and High-Quality Generation</h1>
+        <div class="author-container">
+            <div class="grid-item cmu"><a href="https://gauravparmar.com/">Gaurav Parmar</a></div>
+            <div class="grid-item snap"><a href="https://orpatashnik.github.io/">Or Patashnik</a></div>
+            <div class="grid-item snap"><a href="https://scholar.google.com/citations?user=uD79u6oAAAAJ&hl=en">Daniil Ostashev</a></div>
+            <div class="grid-item snap"><a href="https://wangkua1.github.io/">Kuan-Chieh (Jackson) Wang</a></div>
+            <div class="grid-item snap"><a href="https://kfiraberman.github.io/">Kfir Aberman</a></div>
+        </div>
+        <div class="author-container">
+            <div class="grid-item cmu"><a href="https://www.cs.cmu.edu/~srinivas/">Srinivasa Narasimhan</a></div>
+            <div class="grid-item cmu"><a href="https://www.cs.cmu.edu/~junyanz/">Jun-Yan Zhu</a></div>
+        </div>
+        <br>
+        <div class="affiliation-container">
+            <div class="grid-item cmu"> <p>Carnegie Mellon University</p></div>
+            <div class="grid-item snap"> <p>Snap Research</p></div>
+        </div>
+        
+        <br>
+        <h2>DEMO: Text-to-Image Group Inference with FLUX.1-Schnell</h2>
+        </div>
+        """
+    )
+
+    with gr.Row(scale=1):
+        with gr.Column(scale=1.0):
+            prompt = gr.Textbox(label="Prompt", placeholder="A photo of a dog", lines=4, value="A photo of a dog")
+        
+        with gr.Column(scale=1.0):
+            with gr.Row(elem_id="starting-candidates-row"):
+                gr.Text("Starting Candidates:", container=False, interactive=False, scale=5)
+                starting_candidates = gr.Number(value=default_args.starting_candidates, precision=0, container=False, show_label=False, scale=1)
+            
+            with gr.Row(elem_id="output-group-size-row"):
+                gr.Text("Output Group Size:", container=False, interactive=False, scale=5)
+                output_group_size = gr.Number(value=default_args.output_group_size, precision=0, container=False, show_label=False, scale=1)
+        
+        with gr.Column(scale=1.0):
+            with gr.Accordion("Advanced Options", open=False, elem_id="advanced-options-accordion"):
+                with gr.Row():
+                    gr.Text("Pruning Ratio:", container=False, interactive=False, elem_id="pruning-ratio-label", scale=3)
+                    pruning_ratio = gr.Number(value=default_args.pruning_ratio, precision=2, container=False, show_label=False, scale=1)
+
+                with gr.Row():
+                    gr.Text("Lambda:", container=False, interactive=False, elem_id="lambda-label", scale=5)
+                    lambda_score = gr.Number(value=default_args.lambda_score, precision=1, container=False, show_label=False, scale=1)
+
+                with gr.Row():
+                    gr.Text("Seed:", container=False, interactive=False, elem_id="seed-label", scale=5)
+                    seed = gr.Number(value=42, precision=0, container=False, show_label=False, scale=1)
+                
+                with gr.Row():
+                    gr.Text("Unary:", container=False, interactive=False, elem_id="unary-term-label", scale=2)
+                    unary_term = gr.Dropdown(choices=["clip_text_img"], value=default_args.unary_term, container=False, show_label=False, scale=3)
+                
+                with gr.Row():
+                    gr.Text("Binary:", container=False, interactive=False, elem_id="binary-term-label", scale=2)
+                    binary_term = gr.Dropdown(choices=["diversity_dino", "diversity_clip", "dino_cls_pairwise"], value=default_args.binary_term, 
+                        container=False, show_label=False, scale=3)
+    
+    with gr.Row(scale=1):
+        generate_btn = gr.Button("Generate", variant="primary")
+    
+    with gr.Row(scale=1):
+        output_gallery_group = gr.Gallery(label="Group Inference", show_label=True,elem_id="gallery", columns=4, height="auto")
+    
+    gr.Examples(
+        examples=[
+            ["A photo of a dog", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino"],
+            ["A mountain landscape", 64, 4, 0.5, 1.0, 123, "clip_text_img", "diversity_dino"],
+            ["A cat sleeping", 64, 4, 0.5, 1.0, 456, "clip_text_img", "diversity_dino"],
+            ["A sunset at the beach", 64, 4, 0.5, 1.0, 789, "clip_text_img", "diversity_dino"],
+        ],
+        inputs=[prompt, starting_candidates, output_group_size, pruning_ratio, lambda_score, seed, unary_term, binary_term],
+        outputs=[output_gallery_group],
+        fn=generate_images,
+        cache_examples="lazy",
+        label="Examples"
+    )
+    
+    generate_btn.click(
+        fn=generate_images,
+        inputs=[prompt, starting_candidates, output_group_size, pruning_ratio, lambda_score, seed, unary_term, binary_term],
+        outputs=[output_gallery_group]
+    )
+
+demo.launch()

my_utils/default_values.py

@@ -0,0 +1,81 @@
+DEFAULT_VALUES = {
+    "flux-schnell": {
+        "num_inference_steps": 4,
+        "guidance_scale": 0.0,
+        "starting_candidates": 32,
+        "output_group_size": 4,
+        "pruning_ratio": 0.9,
+        "lambda_score": 1.5,
+        "output_dir": "outputs/flux-schnell",
+        "height": 768,
+        "width": 768,
+        "unary_term": "clip_text_img",
+        "binary_term": "diversity_dino"
+    },
+    "flux-dev": {
+        "num_inference_steps": 20,
+        "guidance_scale": 3.5,
+        "starting_candidates": 128,
+        "output_group_size": 4,
+        "pruning_ratio": 0.5,
+        "lambda_score": 1.5,
+        "output_dir": "outputs/flux-dev",
+        "height": 768,
+        "width": 768,
+        "unary_term": "clip_text_img",
+        "binary_term": "diversity_dino"
+    },
+    "flux-depth": {
+        "num_inference_steps": 20,
+        "guidance_scale": 3.5,
+        "starting_candidates": 128,
+        "output_group_size": 4,
+        "pruning_ratio": 0.5,
+        "lambda_score": 1.5,
+        "output_dir": "outputs/flux-depth",
+        "height": 768,
+        "width": 768,
+        "unary_term": "clip_text_img",
+        "binary_term": "diversity_dino"
+    },
+    "flux-canny": {
+        "num_inference_steps": 20,
+        "guidance_scale": 3.5,
+        "starting_candidates": 128,
+        "output_group_size": 4,
+        "pruning_ratio": 0.5,
+        "lambda_score": 1.5,
+        "output_dir": "outputs/flux-canny",
+        "height": 768,
+        "width": 768,
+        "unary_term": "clip_text_img",
+        "binary_term": "diversity_dino"
+    },
+    "flux-kontext": {
+        "num_inference_steps": 28,
+        "guidance_scale": 3.5,
+        "starting_candidates": 128,
+        "output_group_size": 4,
+        "pruning_ratio": 0.5,
+        "lambda_score": 1.0,
+        "output_dir": "outputs/flux-kontext",
+        "height": 1024,
+        "width": 1024,
+        "unary_term": "clip_text_img",
+        "binary_term": "diversity_dino"
+    }
+}
+
+def apply_defaults(args):
+    model_name = args.model_name
+    
+    if model_name not in DEFAULT_VALUES:
+        raise ValueError(f"Unknown model name: {model_name}. Available models: {list(DEFAULT_VALUES.keys())}")
+    
+    defaults = DEFAULT_VALUES[model_name]
+    
+    for param_name, default_value in defaults.items():
+        if hasattr(args, param_name) and getattr(args, param_name) is None:
+            setattr(args, param_name, default_value)
+    
+    return args 

my_utils/group_inference.py

@@ -0,0 +1,257 @@
+import os, sys
+import math
+import torch
+import spaces
+import numpy as np
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
+
+from my_utils.solvers import gurobi_solver
+
+
+def get_next_size(curr_size, final_size, keep_ratio):
+    """Calculate next size for progressive pruning during denoising.
+    
+    Args:
+        curr_size: Current number of candidates
+        final_size: Target final size
+        keep_ratio: Fraction of candidates to keep at each step
+    """
+    if curr_size < final_size:
+        raise ValueError("Current size is less than the final size!")
+    elif curr_size == final_size:
+        return curr_size
+    else:
+        next_size = math.ceil(curr_size * keep_ratio)
+        return max(next_size, final_size)
+
+
+@torch.no_grad()
+def decode_latent(z, pipe, height, width):
+    """Decode latent tensor to image using VAE decoder.
+    
+    Args:
+        z: Latent tensor to decode
+        pipe: Diffusion pipeline with VAE
+        height: Image height
+        width: Image width
+    """
+    z = pipe._unpack_latents(z, height, width, pipe.vae_scale_factor)
+    z = (z / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
+    z = pipe.vae.decode(z, return_dict=False)[0].clamp(-1,1)
+    return z
+
+
+@torch.no_grad()
+@spaces.GPU(duration=300)
+def run_group_inference(pipe, model_name=None, prompt=None, prompt_2=None, negative_prompt=None, negative_prompt_2=None, 
+        true_cfg_scale=1.0, height=None, width=None, num_inference_steps=28, sigmas=None, guidance_scale=3.5, 
+        l_generator=None, max_sequence_length=512,
+        # group inference arguments
+        unary_score_fn=None, binary_score_fn=None,
+        starting_candidates=None, output_group_size=None, pruning_ratio=None, lambda_score=None,
+        # control arguments
+        control_image=None,
+        # input image for flux-kontext
+        input_image=None,
+        skip_first_cfg=True
+):
+    """Run group inference with progressive pruning for diverse, high-quality image generation.
+    
+    Args:
+        pipe: Diffusion pipeline
+        model_name: Model type (flux-schnell, flux-dev, flux-depth, flux-canny, flux-kontext)
+        prompt: Text prompt for generation
+        unary_score_fn: Function to compute image quality scores
+        binary_score_fn: Function to compute pairwise diversity scores  
+        starting_candidates: Initial number of noise samples
+        output_group_size: Final number of images to generate
+        pruning_ratio: Fraction to prune at each denoising step
+        lambda_score: Weight between quality and diversity terms
+        control_image: Control image for depth/canny models
+        input_image: Input image for flux-kontext editing
+    """
+    if l_generator is None:
+        l_generator = [torch.Generator("cpu").manual_seed(42+_seed) for _seed in range(starting_candidates)]
+
+    # use the default height and width if not provided
+    height = height or pipe.default_sample_size * pipe.vae_scale_factor
+    width = width or pipe.default_sample_size * pipe.vae_scale_factor
+
+    pipe._guidance_scale = guidance_scale
+    pipe._current_timestep = None
+    pipe._interrupt = False
+    pipe._joint_attention_kwargs = {}
+
+    device = pipe._execution_device
+
+    lora_scale = None
+    has_neg_prompt = negative_prompt is not None
+    do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+
+    # 3. Encode prompts
+    prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, prompt_embeds=None, pooled_prompt_embeds=None, device=device, max_sequence_length=max_sequence_length, lora_scale=lora_scale)
+    
+    if do_true_cfg:
+        negative_prompt_embeds, negative_pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=negative_prompt, prompt_2=negative_prompt_2, prompt_embeds=None, pooled_prompt_embeds=None, device=device, max_sequence_length=max_sequence_length, lora_scale=lora_scale)
+
+    # 4. Prepare latent variables
+    if model_name in ["flux-depth", "flux-canny"]:
+        # for control models, the pipe.transformer.config.in_channels is doubled
+        num_channels_latents = pipe.transformer.config.in_channels // 8
+    else:
+        num_channels_latents = pipe.transformer.config.in_channels // 4
+    
+    # Handle different model types
+    image_latents = None
+    image_ids = None
+    if model_name == "flux-kontext":
+        processed_image = pipe.image_processor.preprocess(input_image, height=height, width=width)
+        l_latents = []
+        for _gen in l_generator:
+            latents, img_latents, latent_ids, img_ids = pipe.prepare_latents(
+                processed_image, 1, num_channels_latents, height, width, 
+                prompt_embeds.dtype, device, _gen
+            )
+            l_latents.append(latents)
+        # Use the image_latents and image_ids from the first generator
+        _, image_latents, latent_image_ids, image_ids = pipe.prepare_latents(
+            processed_image, 1, num_channels_latents, height, width, 
+            prompt_embeds.dtype, device, l_generator[0]
+        )
+        # Combine latent_ids with image_ids
+        if image_ids is not None:
+            latent_image_ids = torch.cat([latent_image_ids, image_ids], dim=0)
+    else:
+        # For other models (flux-schnell, flux-dev, flux-depth, flux-canny)
+        l_latents = [pipe.prepare_latents(1, num_channels_latents, height, width, prompt_embeds.dtype, device, _gen)[0] for _gen in l_generator]
+        _, latent_image_ids = pipe.prepare_latents(1, num_channels_latents, height, width, prompt_embeds.dtype, device, l_generator[0])
+    
+    # 4.5. Prepare control image if provided
+    control_latents = None
+    if model_name in ["flux-depth", "flux-canny"]:
+        control_image_processed = pipe.prepare_image(image=control_image, width=width, height=height, batch_size=1, num_images_per_prompt=1, device=device, dtype=pipe.vae.dtype,)
+        if control_image_processed.ndim == 4:
+            control_latents = pipe.vae.encode(control_image_processed).latents
+            control_latents = (control_latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
+            height_control_image, width_control_image = control_latents.shape[2:]
+            control_latents = pipe._pack_latents(control_latents, 1, num_channels_latents, height_control_image, width_control_image)
+
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+    image_seq_len = latent_image_ids.shape[0]
+    mu = calculate_shift(image_seq_len, pipe.scheduler.config.get("base_image_seq_len", 256), pipe.scheduler.config.get("max_image_seq_len", 4096), pipe.scheduler.config.get("base_shift", 0.5), pipe.scheduler.config.get("max_shift", 1.15))
+    timesteps, num_inference_steps = retrieve_timesteps(pipe.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu)
+    num_warmup_steps = max(len(timesteps) - num_inference_steps * pipe.scheduler.order, 0)
+    pipe._num_timesteps = len(timesteps)
+    _dtype = l_latents[0].dtype
+
+    # handle guidance
+    if pipe.transformer.config.guidance_embeds:
+        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(1)
+    else:
+        guidance = None
+    guidance_1 = torch.full([1], 1.0, device=device, dtype=torch.float32).expand(1)
+
+    # 6. Denoising loop
+    with pipe.progress_bar(total=num_inference_steps) as progress_bar:
+        for i, t in enumerate(timesteps):
+            if pipe.interrupt:
+                continue
+            if guidance is not None and skip_first_cfg and i == 0:
+                curr_guidance = guidance_1
+            else:
+                curr_guidance = guidance
+
+            pipe._current_timestep = t
+            timestep = t.expand(1).to(_dtype)
+            # ipdb.set_trace()
+            next_latents = []
+            x0_preds = []
+            # do 1 denoising step
+            for _latent in l_latents:
+                # prepare input for transformer based on model type
+                if model_name in ["flux-depth", "flux-canny"]:
+                    # Control models: concatenate control latents along dim=2
+                    latent_model_input = torch.cat([_latent, control_latents], dim=2)
+                elif model_name == "flux-kontext":
+                    # Kontext model: concatenate image latents along dim=1
+                    latent_model_input = torch.cat([_latent, image_latents], dim=1)
+                else:
+                    # Standard models (flux-schnell, flux-dev): use latents as is
+                    latent_model_input = _latent
+                    
+                noise_pred = pipe.transformer(hidden_states=latent_model_input, timestep=timestep / 1000, guidance=curr_guidance, pooled_projections=pooled_prompt_embeds, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, joint_attention_kwargs=pipe.joint_attention_kwargs, return_dict=False)[0]
+                
+                # For flux-kontext, we need to slice the noise_pred to match the latents size
+                if model_name == "flux-kontext":
+                    noise_pred = noise_pred[:, : _latent.size(1)]
+                
+                if do_true_cfg:
+                    neg_noise_pred = pipe.transformer(hidden_states=latent_model_input, timestep=timestep / 1000, guidance=curr_guidance, pooled_projections=negative_pooled_prompt_embeds, encoder_hidden_states=negative_prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, joint_attention_kwargs=pipe.joint_attention_kwargs, return_dict=False)[0]
+                    if model_name == "flux-kontext":
+                        neg_noise_pred = neg_noise_pred[:, : _latent.size(1)]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                # compute the previous noisy sample x_t -> x_t-1
+                _latent = pipe.scheduler.step(noise_pred, t, _latent, return_dict=False)[0]
+                # the scheduler is not state-less, it maintains a step index that is incremented by one after each step, 
+                # so we need to decrease it here
+                if hasattr(pipe.scheduler, "_step_index"):
+                    pipe.scheduler._step_index -= 1
+                
+                if type(pipe.scheduler) == FlowMatchEulerDiscreteScheduler:
+                    dt = 0.0 - pipe.scheduler.sigmas[i]
+                    x0_pred = _latent + dt * noise_pred
+                else:
+                    raise NotImplementedError("Only Flow Scheduler is supported for now! For other schedulers, you need to manually implement the x0 prediction step.")
+                
+                x0_preds.append(x0_pred)
+                next_latents.append(_latent)
+            
+            if hasattr(pipe.scheduler, "_step_index"):
+                pipe.scheduler._step_index += 1
+            
+            # if the size of next_latents > output_group_size, prune the latents
+            if len(next_latents) > output_group_size:
+                next_size = get_next_size(len(next_latents), output_group_size, 1 - pruning_ratio)
+                print(f"Pruning from {len(next_latents)} to {next_size}")
+                # decode the latents to pixels with tiny-vae
+                l_x0_decoded = [decode_latent(_latent, pipe, height, width) for _latent in x0_preds]
+                # compute the unary and binary scores
+                l_unary_scores = unary_score_fn(l_x0_decoded, target_caption=prompt)
+                M_binary_scores = binary_score_fn(l_x0_decoded) # upper triangular matrix
+                # run with Quadratic Integer Programming sover
+                selected_indices = gurobi_solver(l_unary_scores, M_binary_scores, next_size, lam=lambda_score)
+                l_latents = [next_latents[_i] for _i in selected_indices]
+            else:
+                l_latents = next_latents
+            
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipe.scheduler.order == 0):
+                progress_bar.update()
+
+    pipe._current_timestep = None
+
+    l_images = [pipe._unpack_latents(_latent, height, width, pipe.vae_scale_factor) for _latent in l_latents]
+    l_images = [(latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor for latents in l_images]
+    l_images = [pipe.vae.decode(_image, return_dict=False)[0] for _image in l_images]
+    l_images_tensor = [image.clamp(-1, 1) for image in l_images]  # Keep tensor version for scoring
+    l_images = [pipe.image_processor.postprocess(image, output_type="pil")[0] for image in l_images]
+    
+    # Compute and print final scores
+    print(f"\n=== Final Scores for {len(l_images)} generated images ===")
+    
+    # Compute unary scores
+    final_unary_scores = unary_score_fn(l_images_tensor, target_caption=prompt)
+    print(f"Unary scores (quality): {final_unary_scores}")
+    print(f"Mean unary score: {np.mean(final_unary_scores):.4f}")
+    
+    # Compute binary scores
+    final_binary_scores = binary_score_fn(l_images_tensor)
+    print(f"Binary score matrix (diversity):")
+    print(final_binary_scores)
+    
+    print("=" * 50)
+    
+    pipe.maybe_free_model_hooks()
+    return l_images
+

my_utils/scores.py

@@ -0,0 +1,221 @@
+import functools
+import torch
+import numpy as np
+import torch.nn.functional as F
+import torch.nn as nn
+import torchvision.models as models
+import torchvision.transforms as transforms
+import cv2
+
+from transformers import CLIPProcessor, CLIPModel, AutoModel
+from transformers.models.clip.modeling_clip import _get_vector_norm
+
+
+
+def validate_tensor_list(tensor_list, expected_type=torch.Tensor, min_dims=None, value_range=None, tolerance=0.1):
+    """
+    Validates a list of tensors with specified requirements.
+    
+    Args:
+        tensor_list: List to validate
+        expected_type: Expected type of each element (default: torch.Tensor)
+        min_dims: Minimum number of dimensions each tensor should have
+        value_range: Tuple of (min_val, max_val) for tensor values
+        tolerance: Tolerance for value range checking (default: 0.1)
+    """
+    if not isinstance(tensor_list, list):
+        raise TypeError(f"Input must be a list, got {type(tensor_list)}")
+    
+    if len(tensor_list) == 0:
+        raise ValueError("Input list cannot be empty")
+    
+    for i, item in enumerate(tensor_list):
+        if not isinstance(item, expected_type):
+            raise TypeError(f"List element [{i}] must be {expected_type}, got {type(item)}")
+        
+        if min_dims is not None and len(item.shape) < min_dims:
+            raise ValueError(f"List element [{i}] must have at least {min_dims} dimensions, got shape {item.shape}")
+        
+        if value_range is not None:
+            min_val, max_val = value_range
+            item_min, item_max = item.min().item(), item.max().item()
+            if item_min < (min_val - tolerance) or item_max > (max_val + tolerance):
+                raise ValueError(f"List element [{i}] values must be in range [{min_val}, {max_val}], got range [{item_min:.3f}, {item_max:.3f}]")
+
+
+
+def build_score_fn(name, device="cuda"):
+    """Build scoring functions for image quality and diversity assessment.
+    
+    Args:
+        name: Score function name (clip_text_img, diversity_dino, dino_cls_pairwise, diversity_clip)
+        device: Device to load models on (default: cuda)
+    """
+    d_score_nets = {}
+
+    if name == "clip_text_img":
+        m_clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+        prep_clip = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        score_fn = functools.partial(unary_clip_text_img_t, device=device, m_clip=m_clip, preprocess_clip=prep_clip)
+        d_score_nets["m_clip"] = m_clip
+        d_score_nets["prep_clip"] = prep_clip
+    
+    elif name == "diversity_dino":
+        dino_model = AutoModel.from_pretrained('facebook/dinov2-base').to(device)
+        score_fn = functools.partial(binary_dino_pairwise_t, device=device, dino_model=dino_model)
+        d_score_nets["dino_model"] = dino_model
+    
+    elif name == "dino_cls_pairwise":
+        dino_model = AutoModel.from_pretrained('facebook/dinov2-base').to(device)
+        score_fn = functools.partial(binary_dino_cls_pairwise_t, device=device, dino_model=dino_model)
+        d_score_nets["dino_model"] = dino_model
+    
+    elif name == "diversity_clip":
+        m_clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+        prep_clip = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        score_fn = functools.partial(binary_clip_pairwise_t, device=device, m_clip=m_clip, preprocess_clip=prep_clip)
+        d_score_nets["m_clip"] = m_clip
+        d_score_nets["prep_clip"] = prep_clip
+    
+    else:
+        raise ValueError(f"Invalid score function name: {name}")
+    
+    return score_fn, d_score_nets
+
+
+@torch.no_grad()
+def unary_clip_text_img_t(l_images, device, m_clip, preprocess_clip, target_caption, d_cache=None):
+    """Compute CLIP text-image similarity scores for a list of images.
+    
+    Args:
+        l_images: List of image tensors in range [-1, 1]
+        device: Device for computation
+        m_clip: CLIP model
+        preprocess_clip: CLIP processor
+        target_caption: Text prompt for similarity comparison
+        d_cache: Optional cached text embeddings
+    """
+    # validate input images, l_images should be a list of torch tensors with range [-1, 1]
+    validate_tensor_list(l_images, expected_type=torch.Tensor, min_dims=3, value_range=(-1, 1))
+    
+    _img_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(device)
+    _img_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(device)
+    b_images = torch.cat(l_images, dim=0)
+    b_images = F.interpolate(b_images, size=(224, 224), mode="bilinear", align_corners=False)
+    # re-normalize with clip mean and std
+    b_images = b_images*0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    
+    if d_cache is None:
+        text_encoding = preprocess_clip.tokenizer(target_caption, return_tensors="pt", padding=True).to(device)
+        output = m_clip(pixel_values=b_images, **text_encoding).logits_per_image /m_clip.logit_scale.exp()
+        _score = output.view(-1).cpu().numpy()
+    else:
+        # compute with cached text embeddings
+        vision_outputs = m_clip.vision_model(pixel_values=b_images, output_attentions=False, output_hidden_states=False,
+            interpolate_pos_encoding=False, return_dict=True,)
+        image_embeds = m_clip.visual_projection(vision_outputs[1])
+        image_embeds = image_embeds / _get_vector_norm(image_embeds)
+        text_embeds = d_cache["text_embeds"]
+        _score = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)).t().view(-1).cpu().numpy()
+    
+    return _score
+
+
+@torch.no_grad()
+def binary_dino_pairwise_t(l_images, device, dino_model):
+    """Compute pairwise diversity scores using DINO patch features.
+    
+    Args:
+        l_images: List of image tensors in range [-1, 1]
+        device: Device for computation
+        dino_model: DINO model for feature extraction
+    """
+    # validate input images, l_images should be a list of torch tensors with range [-1, 1]
+    validate_tensor_list(l_images, expected_type=torch.Tensor, min_dims=3, value_range=(-1, 1))
+    
+    b_images = torch.cat(l_images, dim=0)
+    _img_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    _img_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    
+    b_images = F.interpolate(b_images, size=(256, 256), mode="bilinear", align_corners=False)
+    b_images = b_images*0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    all_features = dino_model(pixel_values=b_images).last_hidden_state[:, 1:, :].cpu() # B, 324, 768
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = all_features[i]
+        for j in range(i+1, N):
+            f2 = all_features[j]
+            cos_sim = (1 - F.cosine_similarity(f1, f2, dim=1)).mean().item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix
+
+@torch.no_grad()
+def binary_dino_cls_pairwise_t(l_images, device, dino_model):
+    """Compute pairwise diversity scores using DINO CLS token features.
+    
+    Args:
+        l_images: List of image tensors in range [-1, 1]
+        device: Device for computation
+        dino_model: DINO model for feature extraction
+    """
+    # validate input images, l_images should be a list of torch tensors with range [-1, 1]
+    validate_tensor_list(l_images, expected_type=torch.Tensor, min_dims=3, value_range=(-1, 1))
+    
+    b_images = torch.cat(l_images, dim=0)
+    _img_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    _img_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    
+    b_images = F.interpolate(b_images, size=(256, 256), mode="bilinear", align_corners=False)
+    b_images = b_images*0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    all_features = dino_model(pixel_values=b_images).last_hidden_state[:, 0:1, :].cpu() # B, 1, 768
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = all_features[i]
+        for j in range(i+1, N):
+            f2 = all_features[j]
+            cos_sim = (1 - F.cosine_similarity(f1, f2, dim=1)).mean().item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix
+
+@torch.no_grad()
+def binary_clip_pairwise_t(l_images, device, m_clip, preprocess_clip):
+    """Compute pairwise diversity scores using CLIP image embeddings.
+    
+    Args:
+        l_images: List of image tensors in range [-1, 1]
+        device: Device for computation
+        m_clip: CLIP model
+        preprocess_clip: CLIP processor
+    """
+    # validate input images, l_images should be a list of torch tensors with range [-1, 1]
+    validate_tensor_list(l_images, expected_type=torch.Tensor, min_dims=3, value_range=(-1, 1))
+    
+    _img_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(device)
+    _img_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(device)
+    b_images = torch.cat(l_images, dim=0)
+    b_images = F.interpolate(b_images, size=(224, 224), mode="bilinear", align_corners=False)
+    # re-normalize with clip mean and std
+    b_images = b_images*0.5 + 0.5
+    b_images = (b_images - _img_mean) / _img_std
+    
+    vision_outputs = m_clip.vision_model(pixel_values=b_images, output_attentions=False, output_hidden_states=False,
+            interpolate_pos_encoding=False, return_dict=True,)
+    image_embeds = m_clip.visual_projection(vision_outputs[1])
+    image_embeds = image_embeds / _get_vector_norm(image_embeds)
+    
+    N = len(l_images)
+    score_matrix = np.zeros((N, N))
+    for i in range(N):
+        f1 = image_embeds[i]
+        for j in range(i+1, N):
+            f2 = image_embeds[j]
+            cos_sim = (1 - torch.dot(f1, f2)).item()
+            score_matrix[i, j] = cos_sim
+    return score_matrix

my_utils/solvers.py

@@ -0,0 +1,33 @@
+from gurobipy import Model, GRB, quicksum
+
+
+def gurobi_solver(u, D, n_select, lam=1.0, time_limit=5.0):
+    """Solve quadratic integer programming problem for subset selection with unary and pairwise terms.
+    
+    Args:
+        u: Unary scores for each item
+        D: Pairwise similarity matrix (upper triangular)
+        n_select: Number of items to select
+        lam: Weight for pairwise term (default: 1.0)
+        time_limit: Solver time limit in seconds (default: 5.0)
+    """
+    n = len(u)
+    model = Model()
+    model.Params.LogToConsole = 0
+    model.Params.TimeLimit = time_limit
+    model.Params.OutputFlag = 0
+
+    # Variables: x[i] in {0,1}
+    x = model.addVars(n, vtype=GRB.BINARY, name="x")
+    # Constraint: exactly k items selected
+    model.addConstr(quicksum(x[i] for i in range(n)) == n_select, name="select_k")
+    
+    # Objective: sum of unary + lambda * pairwise
+    linear_part = quicksum(u[i] * x[i] for i in range(n))
+    quadratic_part = quicksum(lam * D[i, j] * x[i] * x[j] for i in range(n) for j in range(i + 1, n))
+
+    model.setObjective(linear_part + quadratic_part, GRB.MAXIMIZE)
+
+    model.optimize()
+    selected_indices = [i for i in range(n) if x[i].X > 0.5]
+    return selected_indices

requirements.txt

@@ -0,0 +1,14 @@
+torch==2.7.1
+torchvision==0.22.1
+torchaudio==2.7.1
+opencv-python
+transformers
+sentencepiece
+protobuf
+accelerate
+diffusers==0.35.1
+gurobipy
+bitsandbytes
+git+https://github.com/openai/CLIP.git
+ipdb
+https://github.com/nunchaku-tech/nunchaku/releases/download/v0.3.1/nunchaku-0.3.1+torch2.7-cp310-cp310-linux_x86_64.whl

styles.css

@@ -0,0 +1,160 @@
+@import url('https://fonts.googleapis.com/css2?family=Varela+Round&display=swap');
+
+::selection {
+    background: rgba(255, 251, 35, 0.58);
+}
+
+body {
+    max-width: 1000px !important;
+    margin: 0 auto !important;
+}
+
+.title_left {
+    padding-top: 0vw !important;
+    filter: none !important; 
+}
+.title_left > h1 {
+    color: #2f2f2f !important;
+    font-family: "Gelasio",Georgia,serif !important;
+    font-weight: normal !important; 
+    font-size: 2.0vw !important; 
+    text-align: center !important; 
+}
+
+.title_left > h2 {
+    color: #2f2f2f !important;
+    font-family: "Gelasio",Georgia,serif !important;
+    font-weight: normal !important; 
+    font-size: 1.5vw !important; 
+    text-align: center !important; 
+}
+
+.author-container {
+    color: #2f2f2f;
+    font-family: Gelasio,"Avenir Next",Helvetica,sans-serif;
+    font-weight: normal;
+    font-size: 1vw;
+    padding-top: 0.2vw;
+    justify-items: center;
+    justify-content: center;
+    display: grid;
+    grid-template-columns: auto auto auto auto auto;
+}
+
+.affiliation-container {
+    color: #2f2f2f;
+    font-family: Gelasio,"Avenir Next",Helvetica,sans-serif;
+    font-weight: normal;
+    font-size: 1vw;
+    padding-top: 0.2vw;
+    justify-items: center;
+    justify-content: center;
+    display: grid;
+    grid-template-columns: auto auto auto auto auto;
+}
+
+.grid-item {
+    text-align: center;
+    padding-right: 0.7vw;
+    padding-left: 0.7vw;
+}
+
+.grid-item > a {
+    color: #2f2f2f;
+    text-decoration: underline;
+    text-underline-offset: 3px;
+}
+
+.grid-item.cmu > a {
+    text-decoration-color: rgba(196, 18, 48, 0.2)
+}
+
+.grid-item.cmu > p::before {
+    content: "";
+    display: inline-block;
+    width: 12px;
+    height: 12px;
+    background-color: rgba(196, 18, 48, 0.6);
+    margin-right: 8px;
+    vertical-align: middle;
+}
+
+.grid-item.snap > a {
+    text-decoration-color: rgba(255,252,0, 0.4)
+}
+
+.grid-item.snap > p::before {
+    content: "";
+    display: inline-block;
+    width: 12px;
+    height: 12px;
+    background-color: rgba(255, 252, 0, 0.6);
+    margin-right: 8px;
+    vertical-align: middle;
+}
+
+.grid-item.cmu > p,
+.grid-item.snap > p {
+    color: #2f2f2f;
+}
+
+.column {
+    min-width: min(100px, 100%) !important;
+}
+
+.block {
+    min-width: min(100px, 100%) !important;
+    background: transparent !important;
+    border: none !important;
+}
+
+.gr-box, .gr-form, .gr-panel {
+    background: transparent !important;
+    border: none !important;
+}
+
+.gr-row, .gr-column {
+    background: transparent !important;
+}
+
+.gr-textbox, .gr-number, .gr-slider, .gr-dropdown {
+    background: rgba(255, 255, 255, 0.1) !important;
+    border: 1px solid rgba(255, 255, 255, 0.2) !important;
+    backdrop-filter: blur(10px) !important;
+}
+
+.gr-button {
+    background: rgba(255, 255, 255, 0.15) !important;
+    border: 1px solid rgba(255, 255, 255, 0.3) !important;
+    backdrop-filter: blur(10px) !important;
+}
+
+.gr-accordion {
+    background: rgba(255, 255, 255, 0.05) !important;
+    border: 1px solid rgba(255, 255, 255, 0.1) !important;
+    backdrop-filter: blur(5px) !important;
+}
+
+.gr-gallery {
+    background: rgba(255, 255, 255, 0.05) !important;
+    border: 1px solid rgba(255, 255, 255, 0.1) !important;
+    backdrop-filter: blur(5px) !important;
+}
+
+#starting-candidates-row > #component-7 {
+    border: none !important;
+    /* font-family: "Varela Round" !important;
+    font-weight: 500 !important;  */
+}
+
+#output-group-size-row  >#component-10{
+    border: none !important;
+    /* font-family: "Varela Round" !important; */
+    /* font-weight: 500 !important;  */
+}
+
+#pruning-ratio-label, #lambda-label, #seed-label, #unary-term-label, #binary-term-label {
+    border: none !important;
+    /* font-family: "Varela Round" !important;
+    font-weight: 500 !important;  */
+}