derekl35/FLUX-Quantization-Challenge
0
1import torch2import gradio as gr3from diffusers import FluxPipeline, FluxTransformer2DModel4from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig5from transformers import T5EncoderModel6from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig7import gc8import random9from PIL import Image10import os11import time12import spaces13 14DEVICE = "cuda" if torch.cuda.is_available() else "cpu"15print(f"Using device: {DEVICE}")16 17DEFAULT_HEIGHT = 102418DEFAULT_WIDTH = 102419DEFAULT_GUIDANCE_SCALE = 3.520DEFAULT_NUM_INFERENCE_STEPS = 5021DEFAULT_MAX_SEQUENCE_LENGTH = 51222GENERATION_SEED = 0 # could use a random number generator to set this, for more variety23 24def clear_gpu_memory(*args):25 allocated_before = torch.cuda.memory_allocated(0) / 1024**3 if DEVICE == "cuda" else 026 reserved_before = torch.cuda.memory_reserved(0) / 1024**3 if DEVICE == "cuda" else 027 print(f"Before clearing: Allocated={allocated_before:.2f} GB, Reserved={reserved_before:.2f} GB")28 29 deleted_types = []30 for arg in args:31 if arg is not None:32 deleted_types.append(str(type(arg)))33 del arg34 35 if deleted_types:36 print(f"Deleted objects of types: {', '.join(deleted_types)}")37 else:38 print("No objects passed to clear_gpu_memory.")39 40 gc.collect()41 if DEVICE == "cuda":42 torch.cuda.empty_cache()43 44 allocated_after = torch.cuda.memory_allocated(0) / 1024**3 if DEVICE == "cuda" else 045 reserved_after = torch.cuda.memory_reserved(0) / 1024**3 if DEVICE == "cuda" else 046 print(f"After clearing: Allocated={allocated_after:.2f} GB, Reserved={reserved_after:.2f} GB")47 print("-" * 20)48 49CACHED_PIPES = {}50def load_bf16_pipeline():51 """Loads the original FLUX.1-dev pipeline in BF16 precision."""52 print("Loading BF16 pipeline...")53 MODEL_ID = "black-forest-labs/FLUX.1-dev"54 if MODEL_ID in CACHED_PIPES:55 return CACHED_PIPES[MODEL_ID]56 start_time = time.time()57 try:58 pipe = FluxPipeline.from_pretrained(59 MODEL_ID,60 torch_dtype=torch.bfloat1661 )62 pipe.to(DEVICE)63 # pipe.enable_model_cpu_offload()64 end_time = time.time()65 mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 066 print(f"BF16 Pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")67 # CACHED_PIPES[MODEL_ID] = pipe68 return pipe69 except Exception as e:70 print(f"Error loading BF16 pipeline: {e}")71 raise # Re-raise exception to be caught in generate_images72 73def load_bnb_8bit_pipeline():74 """Loads the FLUX.1-dev pipeline with 8-bit quantized components."""75 print("Loading 8-bit BNB pipeline...")76 MODEL_ID = "derekl35/FLUX.1-dev-bnb-8bit"77 if MODEL_ID in CACHED_PIPES:78 return CACHED_PIPES[MODEL_ID]79 start_time = time.time()80 try:81 pipe = FluxPipeline.from_pretrained(82 MODEL_ID,83 torch_dtype=torch.bfloat1684 )85 pipe.to(DEVICE)86 # pipe.enable_model_cpu_offload()87 end_time = time.time()88 mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 089 print(f"8-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")90 CACHED_PIPES[MODEL_ID] = pipe91 return pipe92 except Exception as e:93 print(f"Error loading 8-bit BNB pipeline: {e}")94 raise95 96def load_bnb_4bit_pipeline():97 """Loads the FLUX.1-dev pipeline with 4-bit quantized components."""98 print("Loading 4-bit BNB pipeline...")99 MODEL_ID = "derekl35/FLUX.1-dev-nf4"100 if MODEL_ID in CACHED_PIPES:101 return CACHED_PIPES[MODEL_ID]102 start_time = time.time()103 try:104 pipe = FluxPipeline.from_pretrained(105 MODEL_ID,106 torch_dtype=torch.bfloat16107 )108 pipe.to(DEVICE)109 # pipe.enable_model_cpu_offload()110 end_time = time.time()111 mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0112 print(f"4-bit BNB pipeline loaded in {end_time - start_time:.2f}s. Memory reserved: {mem_reserved:.2f} GB")113 CACHED_PIPES[MODEL_ID] = pipe114 return pipe115 except Exception as e:116 print(f"4-bit BNB pipeline: {e}")117 raise118 119@spaces.GPU(duration=240)120def generate_images(prompt, quantization_choice, progress=gr.Progress(track_tqdm=True)):121 """Loads original and selected quantized model, generates one image each, clears memory, shuffles results."""122 if not prompt:123 return None, {}, gr.update(value="Please enter a prompt.", interactive=False), gr.update(choices=[], value=None)124 125 if not quantization_choice:126 # Return updates for all outputs to clear them or show warning127 return None, {}, gr.update(value="Please select a quantization method.", interactive=False), gr.update(choices=[], value=None)128 129 # Determine which quantized model to load130 if quantization_choice == "8-bit":131 quantized_load_func = load_bnb_8bit_pipeline132 quantized_label = "Quantized (8-bit)"133 elif quantization_choice == "4-bit":134 quantized_load_func = load_bnb_4bit_pipeline135 quantized_label = "Quantized (4-bit)"136 else:137 # Should not happen with Radio choices, but good practice138 return None, {}, gr.update(value="Invalid quantization choice.", interactive=False), gr.update(choices=[], value=None)139 140 model_configs = [141 ("Original", load_bf16_pipeline),142 (quantized_label, quantized_load_func), # Use the specific label here143 ]144 145 results = []146 pipe_kwargs = {147 "prompt": prompt,148 "height": DEFAULT_HEIGHT,149 "width": DEFAULT_WIDTH,150 "guidance_scale": DEFAULT_GUIDANCE_SCALE,151 "num_inference_steps": DEFAULT_NUM_INFERENCE_STEPS,152 "max_sequence_length": DEFAULT_MAX_SEQUENCE_LENGTH,153 }154 155 current_pipe = None # Keep track of the current pipe for cleanup156 157 for i, (label, load_func) in enumerate(model_configs):158 progress(i / len(model_configs), desc=f"Loading {label} model...")159 print(f"\n--- Loading {label} Model ---")160 load_start_time = time.time()161 try:162 # Ensure previous pipe is cleared *before* loading the next163 # if current_pipe:164 # print(f"--- Clearing memory before loading {label} Model ---")165 # clear_gpu_memory(current_pipe)166 # current_pipe = None167 168 current_pipe = load_func()169 load_end_time = time.time()170 print(f"{label} model loaded in {load_end_time - load_start_time:.2f} seconds.")171 172 progress((i + 0.5) / len(model_configs), desc=f"Generating with {label} model...")173 print(f"--- Generating with {label} Model ---")174 gen_start_time = time.time()175 image_list = current_pipe(**pipe_kwargs, generator=torch.manual_seed(GENERATION_SEED)).images176 image = image_list[0]177 gen_end_time = time.time()178 results.append({"label": label, "image": image})179 print(f"--- Finished Generation with {label} Model in {gen_end_time - gen_start_time:.2f} seconds ---")180 mem_reserved = torch.cuda.memory_reserved(0)/1024**3 if DEVICE == "cuda" else 0181 print(f"Memory reserved: {mem_reserved:.2f} GB")182 183 except Exception as e:184 print(f"Error during {label} model processing: {e}")185 # Attempt cleanup186 if current_pipe:187 print(f"--- Clearing memory after error with {label} Model ---")188 clear_gpu_memory(current_pipe)189 current_pipe = None190 # Return error state to Gradio - update all outputs191 return None, {}, gr.update(value=f"Error processing {label} model: {e}", interactive=False), gr.update(choices=[], value=None)192 193 # No finally block needed here, cleanup happens before next load or after loop194 195 # Final cleanup after the loop finishes successfully196 # if current_pipe:197 # print(f"--- Clearing memory after last model ({label}) ---")198 # clear_gpu_memory(current_pipe)199 # current_pipe = None200 201 if len(results) != len(model_configs):202 print("Generation did not complete for all models.")203 # Update all outputs204 return None, {}, gr.update(value="Failed to generate images for all model types.", interactive=False), gr.update(choices=[], value=None)205 206 # Shuffle the results for display207 shuffled_results = results.copy()208 random.shuffle(shuffled_results)209 210 # Create the gallery data: [(image, caption), (image, caption)]211 shuffled_data_for_gallery = [(res["image"], f"Image {i+1}") for i, res in enumerate(shuffled_results)]212 213 # Create the mapping: display_index -> correct_label (e.g., {0: 'Original', 1: 'Quantized (8-bit)'})214 correct_mapping = {i: res["label"] for i, res in enumerate(shuffled_results)}215 print("Correct mapping (hidden):", correct_mapping)216 217 guess_radio_update = gr.update(choices=["Image 1", "Image 2"], value=None, interactive=True)218 219 # Return shuffled images, the correct mapping state, status message, and update the guess radio220 return shuffled_data_for_gallery, correct_mapping, gr.update(value="Generation complete! Make your guess.", interactive=False), guess_radio_update221 222 223# --- Guess Verification Function ---224def check_guess(user_guess, correct_mapping_state):225 """Compares the user's guess with the correct mapping stored in the state."""226 227 if not isinstance(correct_mapping_state, dict) or not correct_mapping_state:228 return "Please generate images first (state is empty or invalid)."229 230 if user_guess is None:231 return "Please select which image you think is quantized."232 233 # Find which display index (0 or 1) corresponds to the quantized image234 quantized_image_index = -1235 quantized_label_actual = ""236 for index, label in correct_mapping_state.items():237 if "Quantized" in label: # Check if the label indicates quantization238 quantized_image_index = index239 quantized_label_actual = label # Store the full label e.g. "Quantized (8-bit)"240 break241 242 if quantized_image_index == -1:243 # This shouldn't happen if generation was successful244 return "Error: Could not find the quantized image in the mapping data."245 246 # Determine what the user *should* have selected based on the index247 correct_guess_label = f"Image {quantized_image_index + 1}" # "Image 1" or "Image 2"248 249 if user_guess == correct_guess_label:250 feedback = f"Correct! {correct_guess_label} used the {quantized_label_actual} model."251 else:252 feedback = f"Incorrect. The quantized image ({quantized_label_actual}) was {correct_guess_label}."253 254 return feedback255 256 257with gr.Blocks(title="FLUX Quantization Challenge", theme=gr.themes.Soft()) as demo:258 gr.Markdown("# FLUX Model Quantization Challenge")259 gr.Markdown(260 "Compare the original FLUX.1-dev (BF16) model against a quantized version (4-bit or 8-bit). "261 "Enter a prompt, choose the quantization method, and generate two images. "262 "The images will be shuffled. Can you guess which one used quantization?"263 )264 265 with gr.Row():266 prompt_input = gr.Textbox(label="Enter Prompt", placeholder="e.g., A photorealistic portrait of an astronaut on Mars", scale=3)267 quantization_choice_radio = gr.Radio(268 choices=["8-bit", "4-bit"],269 label="Select Quantization",270 value="8-bit", # Default choice271 scale=1272 )273 generate_button = gr.Button("Generate & Compare", variant="primary", scale=1)274 275 output_gallery = gr.Gallery(276 label="Generated Images (Original vs. Quantized)",277 columns=2,278 height=512,279 object_fit="contain",280 allow_preview=True,281 show_label=True, # Shows "Image 1", "Image 2" captions we provide282 )283 284 gr.Markdown("### Which image used the selected quantization method?")285 with gr.Row():286 # Centered guess radio and submit button287 with gr.Column(scale=1): # Dummy column for spacing288 pass289 with gr.Column(scale=2): # Column for the radio button290 guess_radio = gr.Radio(291 choices=[],292 label="Your Guess",293 info="Select the image you believe was generated with the quantized model.",294 interactive=False # Disabled until images are generated295 )296 with gr.Column(scale=1): # Column for the button297 submit_guess_button = gr.Button("Submit Guess")298 with gr.Column(scale=1): # Dummy column for spacing299 pass300 301 feedback_box = gr.Textbox(label="Feedback", interactive=False, lines=1)302 303 # Hidden state to store the correct mapping after shuffling304 # e.g., {0: 'Original', 1: 'Quantized (8-bit)'} or {0: 'Quantized (4-bit)', 1: 'Original'}305 correct_mapping_state = gr.State({})306 307 generate_button.click(308 fn=generate_images,309 inputs=[prompt_input, quantization_choice_radio],310 outputs=[output_gallery, correct_mapping_state, feedback_box, guess_radio]311 ).then(312 lambda: "", # Clear feedback box on new generation313 outputs=[feedback_box]314 )315 316 317 submit_guess_button.click(318 fn=check_guess,319 inputs=[guess_radio, correct_mapping_state], # Pass the selected guess and the state320 outputs=[feedback_box]321 )322 323if __name__ == "__main__":324 # queue()325 # demo.queue().launch() # Set share=True to create public link if needed326 demo.launch()