import os import gradio as gr import numpy as np import random import spaces import torch from diffusers import Flux2KleinPipeline from huggingface_hub import login, hf_hub_download from safetensors.torch import load_file try: from color_matcher import ColorMatcher except ImportError: import subprocess subprocess.check_call(["pip", "install", "color-matcher"]) from color_matcher import ColorMatcher dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 # Login to HuggingFace for private repos if token is available if "HF_TOKEN_LOGIN" in os.environ: login(token=os.environ["HF_TOKEN_LOGIN"]) # Model repository IDs REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-4B" REPO_ID_ALBEDO_LORA = "NightRaven109/kleinalbedo4B5ksteps" # Load 4B Distilled model with Albedo LoRA print("Loading 4B Distilled model...") pipe = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype) print("Loading Albedo LoRA...") # Download and load LoRA weights manually lora_path = hf_hub_download(repo_id=REPO_ID_ALBEDO_LORA, filename="klein_albedo_5k_diffprompt.safetensors") lora_state_dict = load_file(lora_path) # Merge LoRA weights into model with proper key conversion model_state_dict = pipe.transformer.state_dict() lora_scale = 1.0 # Get hidden size from model weights hidden_size = model_state_dict["transformer_blocks.0.attn.to_q.weight"].shape[0] print(f"Model hidden size: {hidden_size}") # Check LoRA dimensions sample_qkv_lora_b = lora_state_dict.get("diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight") if sample_qkv_lora_b is not None: lora_qkv_size = sample_qkv_lora_b.shape[0] lora_hidden_size = lora_qkv_size // 3 print(f"LoRA QKV size: {lora_qkv_size}, LoRA hidden size: {lora_hidden_size}") matched_keys = 0 unmatched_keys = [] # Store LoRA deltas for runtime adjustment lora_deltas = {} current_lora_scale = [lora_scale] # Use list to allow modification in nested function def apply_lora(model_state_dict, base_key, lora_a, lora_b, scale, store_delta=True): """Apply LoRA delta to model weights.""" if base_key in model_state_dict: delta = (lora_b @ lora_a) if store_delta: lora_deltas[base_key] = delta.clone() model_state_dict[base_key] = model_state_dict[base_key] + (delta * scale).to(model_state_dict[base_key].dtype) return True return False for key in lora_state_dict.keys(): if not key.endswith(".lora_A.weight"): continue lora_b_key = key.replace(".lora_A.weight", ".lora_B.weight") if lora_b_key not in lora_state_dict: continue lora_a = lora_state_dict[key].to(dtype=dtype) lora_b = lora_state_dict[lora_b_key].to(dtype=dtype) # Remove prefix and get base name base = key.replace(".lora_A.weight", "") if base.startswith("diffusion_model."): base = base[len("diffusion_model."):] applied = False # Double blocks - image attention QKV (split into Q, K, V) if ".img_attn.qkv" in base: block_prefix = base.replace(".img_attn.qkv", "").replace("double_blocks.", "transformer_blocks.") # Split lora_b into Q, K, V parts q_b = lora_b[0:hidden_size, :] k_b = lora_b[hidden_size:2*hidden_size, :] v_b = lora_b[2*hidden_size:3*hidden_size, :] applied = apply_lora(model_state_dict, f"{block_prefix}.attn.to_q.weight", lora_a, q_b, lora_scale) applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.to_k.weight", lora_a, k_b, lora_scale) applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.to_v.weight", lora_a, v_b, lora_scale) # Double blocks - text attention QKV (split into add_q, add_k, add_v) elif ".txt_attn.qkv" in base: block_prefix = base.replace(".txt_attn.qkv", "").replace("double_blocks.", "transformer_blocks.") # Split lora_b into Q, K, V parts q_b = lora_b[0:hidden_size, :] k_b = lora_b[hidden_size:2*hidden_size, :] v_b = lora_b[2*hidden_size:3*hidden_size, :] applied = apply_lora(model_state_dict, f"{block_prefix}.attn.add_q_proj.weight", lora_a, q_b, lora_scale) applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.add_k_proj.weight", lora_a, k_b, lora_scale) applied |= apply_lora(model_state_dict, f"{block_prefix}.attn.add_v_proj.weight", lora_a, v_b, lora_scale) # Double blocks - image attention output projection elif ".img_attn.proj" in base: new_key = base.replace("double_blocks.", "transformer_blocks.").replace(".img_attn.proj", ".attn.to_out.0") + ".weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Double blocks - text attention output projection elif ".txt_attn.proj" in base: new_key = base.replace("double_blocks.", "transformer_blocks.").replace(".txt_attn.proj", ".attn.to_add_out") + ".weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Single blocks - linear1 maps to to_qkv_mlp_proj elif "single_blocks." in base and ".linear1" in base: new_key = base.replace("single_blocks.", "single_transformer_blocks.").replace(".linear1", ".attn.to_qkv_mlp_proj") + ".weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Single blocks - linear2 maps to to_out elif "single_blocks." in base and ".linear2" in base: new_key = base.replace("single_blocks.", "single_transformer_blocks.").replace(".linear2", ".attn.to_out") + ".weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Modulation layers elif "double_stream_modulation_img.lin" in base: new_key = "double_stream_modulation_img.linear.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) elif "double_stream_modulation_txt.lin" in base: new_key = "double_stream_modulation_txt.linear.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) elif "single_stream_modulation.lin" in base: new_key = "single_stream_modulation.linear.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Embedders elif "img_in" in base: new_key = "x_embedder.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) elif "txt_in" in base: new_key = "context_embedder.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Time embedding elif "time_in.in_layer" in base: new_key = "time_guidance_embed.timestep_embedder.linear_1.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) elif "time_in.out_layer" in base: new_key = "time_guidance_embed.timestep_embedder.linear_2.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) # Final layer elif "final_layer.linear" in base: new_key = "proj_out.weight" applied = apply_lora(model_state_dict, new_key, lora_a, lora_b, lora_scale) if applied: matched_keys += 1 else: unmatched_keys.append(base) print(f"LoRA keys matched: {matched_keys}") print(f"LoRA keys unmatched: {len(unmatched_keys)}") if unmatched_keys: print("Unmatched keys:", unmatched_keys) pipe.transformer.load_state_dict(model_state_dict) print("LoRA merged successfully!") print(f"Stored {len(lora_deltas)} LoRA deltas for runtime adjustment") def update_lora_scale(new_scale): """Update LoRA scale at runtime by adjusting weights.""" global current_lora_scale old_scale = current_lora_scale[0] if new_scale == old_scale: return scale_diff = new_scale - old_scale state_dict = pipe.transformer.state_dict() for key, delta in lora_deltas.items(): if key in state_dict: target_device = state_dict[key].device delta_on_device = delta.to(device=target_device, dtype=state_dict[key].dtype) state_dict[key] = state_dict[key] + (delta_on_device * scale_diff) pipe.transformer.load_state_dict(state_dict) current_lora_scale[0] = new_scale print(f"LoRA scale updated: {old_scale} -> {new_scale}") pipe.to("cuda") def update_dimensions_from_image(image_list): """Update width/height sliders based on uploaded image aspect ratio. Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8.""" if image_list is None or len(image_list) == 0: return 1024, 1024 # Default dimensions # Get the first image to determine dimensions img = image_list[0][0] # Gallery returns list of tuples (image, caption) img_width, img_height = img.size aspect_ratio = img_width / img_height if aspect_ratio >= 1: # Landscape or square new_width = 1024 new_height = int(1024 / aspect_ratio) else: # Portrait new_height = 1024 new_width = int(1024 * aspect_ratio) # Round to nearest multiple of 8 new_width = round(new_width / 8) * 8 new_height = round(new_height / 8) * 8 # Ensure within valid range (minimum 256, maximum 1024) new_width = max(256, min(1024, new_width)) new_height = max(256, min(1024, new_height)) return new_width, new_height def reinhard_color_match(target, reference, strength=1.0): """Apply Reinhard color transfer from reference image to target image. Both inputs and output are PIL Images.""" cm = ColorMatcher() target_np = np.array(target).astype(np.float32) / 255.0 reference_np = np.array(reference).astype(np.float32) / 255.0 matched_np = cm.transfer(src=target_np, ref=reference_np, method='reinhard') if strength < 1.0: matched_np = target_np + strength * (matched_np - target_np) matched_np = np.clip(matched_np * 255, 0, 255).astype(np.uint8) from PIL import Image return Image.fromarray(matched_np) PROMPT = "Basecolor, no shadows, flat lighting, keep color" @spaces.GPU(duration=85) def infer(input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=1, guidance_scale=1.0, lora_weight=1.0, color_match_enabled=True, color_match_strength=1.0, progress=gr.Progress(track_tqdm=True)): # Update LoRA scale if changed update_lora_scale(lora_weight) if randomize_seed: seed = random.randint(0, MAX_SEED) # Prepare image list (convert None or empty gallery to None) image_list = None if input_images is not None and len(input_images) > 0: image_list = [] for item in input_images: image_list.append(item[0]) progress(0.2, desc="Generating image...") generator = torch.Generator(device=device).manual_seed(seed) pipe_kwargs = { "prompt": PROMPT, "height": height, "width": width, "num_inference_steps": num_inference_steps, "guidance_scale": guidance_scale, "generator": generator, } # Add images if provided if image_list is not None: pipe_kwargs["image"] = image_list output_image = pipe(**pipe_kwargs).images[0] # Apply Reinhard color match to transfer input colors to output if color_match_enabled and image_list is not None and color_match_strength > 0: progress(0.8, desc="Applying Reinhard color match...") output_image = reinhard_color_match(output_image, image_list[0], strength=color_match_strength) # Create comparison tuple (input, output) for slider input_for_compare = image_list[0] if image_list else output_image comparison = (input_for_compare, output_image) return output_image, comparison, seed css = """ #col-container { margin: 0 auto; max-width: 1200px; } .gallery-container img{ object-fit: contain; } """ with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown(f"""# FLUX.2 [Klein] - 4B (Apache 2.0) FLUX.2 [klein] is a fast, unified image generation and editing model designed for fast inference [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)], [[blog](https://bfl.ai/blog/flux-2)] """) with gr.Row(): with gr.Column(): run_button = gr.Button("Run", scale=1) input_images = gr.Gallery( label="Input Image(s)", type="pil", columns=3, rows=1, ) with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024, ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=8, value=1024, ) with gr.Row(): num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=100, step=1, value=1, ) guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0, ) lora_weight = gr.Slider( label="LoRA Weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, ) color_match_enabled = gr.Checkbox(label="Reinhard Color Match", value=True) color_match_strength = gr.Slider( label="Color Match Strength", minimum=0.0, maximum=1.0, step=0.05, value=1.0, ) with gr.Column(): result = gr.Image(label="Result", show_label=False) comparison_slider = gr.ImageSlider( label="Before / After Comparison", type="pil", ) # Auto-update dimensions when images are uploaded input_images.upload( fn=update_dimensions_from_image, inputs=[input_images], outputs=[width, height] ) gr.on( triggers=[run_button.click], fn=infer, inputs=[input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, lora_weight, color_match_enabled, color_match_strength], outputs=[result, comparison_slider, seed] ) demo.launch(css=css)