import os import spaces import time import gradio as gr import torch import functools import numpy as np 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 my_utils.group_inference import run_group_inference from my_utils.default_values import apply_defaults 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() 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-kontext", 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=512, width=512, ) 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=200) def generate_images(prompt, starting_candidates, output_group_size, pruning_ratio, lambda_score, seed, unary_term, binary_term, input_image=None, 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-kontext", "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, } inference_args["input_image"] = input_image print(f"pruning ratio is: {pruning_ratio}") # Run group inference t_start = time.time() output_group = run_group_inference(pipe, **inference_args) t_end = time.time() print(f"Time taken for group inference: {t_end - t_start} seconds") 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( """

Scaling Group Inference for Diverse and High-Quality Generation

Gaurav Parmar
Or Patashnik
Daniil Ostashev
Kuan-Chieh (Jackson) Wang
Kfir Aberman
Srinivasa Narasimhan
Jun-Yan Zhu

Carnegie Mellon University

Snap Research


DEMO: Text-to-Image Group Inference with FLUX.1-Schnell

""" ) with gr.Row(scale=1): with gr.Column(scale=1.0): prompt_placeholder = "Cat is playing outside in nature." prompt_default = "Cat is playing outside in nature." prompt = gr.Textbox(label="Prompt", placeholder=prompt_placeholder, lines=4, value=prompt_default) input_image = gr.Image(label="Input Image", type="pil", sources=["upload"]) 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) # Instructions for users gr.HTML( """
Tips:
""" ) 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") generate_btn.click( fn=generate_images, inputs=[prompt, starting_candidates, output_group_size, pruning_ratio, lambda_score, seed, unary_term, binary_term, input_image], 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()