Adding app.py

app.py

@@ -0,0 +1,298 @@
+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 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")
+
+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(
+        """
+        <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_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)
+    
+    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]
+    )
+
+demo.launch()