4bit and benchmarked

app.py

@@ -1,5 +1,6 @@
 import os
 import spaces
+import time
 import gradio as gr
 import torch
 import functools
@@ -8,11 +9,16 @@ 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 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
+
 import argparse
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")
+precision = get_precision()
+transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"nunchaku-tech/nunchaku-flux.1-schnell/svdq-{precision}_r32-flux.1-schnell.safetensors")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1").to("cuda")
 
 m_clip = CLIPModel.from_pretrained("multimodalart/clip-vit-base-patch32").to("cuda")
@@ -178,7 +184,10 @@ def generate_images(prompt, starting_candidates, output_group_size, pruning_rati
     }
     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
 
 

my_utils/group_inference.py

@@ -1,4 +1,4 @@
-import os, sys
+import os, sys, time
 import math
 import torch
 import spaces
@@ -221,7 +221,10 @@ def run_group_inference(pipe, model_name=None, prompt=None, prompt_2=None, negat
                 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