adding utils:

app.py

@@ -9,22 +9,13 @@ import torch.nn.functional as F
 from diffusers import FluxPipeline, AutoencoderTiny, FluxKontextPipeline
 from transformers import CLIPProcessor, CLIPModel, AutoModel
 from transformers.models.clip.modeling_clip import _get_vector_norm
-from nunchaku import NunchakuFluxTransformer2dModel
-from nunchaku.utils import get_precision
 from my_utils.group_inference import run_group_inference
 from my_utils.default_values import apply_defaults
-from diffusers.hooks import apply_group_offloading
-from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, FluxTransformer2DModel, FluxPipeline
-from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel
-
 import argparse
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
-pipe.enable_model_cpu_offload()
-
-
-# pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
+# pipe.enable_model_cpu_offload()
 
 m_clip = CLIPModel.from_pretrained("multimodalart/clip-vit-base-patch32").to("cuda")
 prep_clip = CLIPProcessor.from_pretrained("multimodalart/clip-vit-base-patch32")
@@ -283,6 +274,22 @@ with gr.Blocks(css=custom_css, js=js_func, theme=gr.themes.Soft(), elem_id="main
                     binary_term = gr.Dropdown(choices=["diversity_dino", "diversity_clip", "dino_cls_pairwise"], value=default_args.binary_term, 
                         container=False, show_label=False, scale=3)
     
+
+    # Instructions for users
+    gr.HTML(
+        """
+        <div style="margin: 15px 0; padding: 10px; background-color: #f0f0f0; border-radius: 5px; font-size: 14px;">
+            <strong>Tips:</strong>
+            <ul style="margin: 5px 0; padding-left: 20px;">
+                <li>Try out the (cached) examples below first! </li>
+                <li>Higher lambda → more diverse outputs (no added runtime cost)</li>
+                <li>Lower lambda → improved quality and text-adherence (no added runtime cost)</li>
+                <li>More starting candidates → better quality and diversity (slower runtime)</li>
+            </ul>
+        </div>
+        """
+    )
+    
     with gr.Row(scale=1):
         generate_btn = gr.Button("Generate", variant="primary")
     
@@ -295,4 +302,19 @@ with gr.Blocks(css=custom_css, js=js_func, theme=gr.themes.Soft(), elem_id="main
         outputs=[output_gallery_group]
     )
 
+    gr.Examples(
+        examples=[
+            ["Cat is sitting in a cafe and working on his laptop.", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino", "assets/cat.png"],
+            ["Cat is playing outside in nature.", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino", "assets/cat.png"],
+            ["Cat is drinking a glass of milk.", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino", "assets/cat.png"],
+            ["Cat is an astronaut landing on the moon.", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino", "assets/cat.png"],
+            ["Cat is surfing in the ocean.", 64, 4, 0.5, 1.0, 42, "clip_text_img", "diversity_dino", "assets/cat.png"],
+        ],
+        inputs=[prompt, starting_candidates, output_group_size, pruning_ratio, lambda_score, seed, unary_term, binary_term, input_image],
+        outputs=[output_gallery_group],
+        fn=generate_images,
+        cache_examples=True,
+        label="Examples"
+    )
+
 demo.launch()

my_utils/default_values.py

@@ -0,0 +1,29 @@
+DEFAULT_VALUES = {
+    "flux-kontext": {
+        "num_inference_steps": 28,
+        "guidance_scale": 3.5,
+        "starting_candidates": 32,
+        "output_group_size": 4,
+        "pruning_ratio": 0.5,
+        "lambda_score": 1.0,
+        "output_dir": "outputs/flux-kontext",
+        "height": 512,
+        "width": 512,
+        "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,260 @@
+import os, sys, time
+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
+                t_start = time.time()
+                selected_indices = gurobi_solver(l_unary_scores, M_binary_scores, next_size, lam=lambda_score)
+                t_end = time.time()
+                print(f"Time taken for QIP: {t_end - t_start} seconds")
+                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