MagnaSC/FLUX.2-klein-9B
0
1import os2import subprocess3import sys4import io5import gradio as gr6import numpy as np7import random8import spaces9import torch10from diffusers import Flux2KleinPipeline11import requests12from PIL import Image13import json14import base6415from huggingface_hub import InferenceClient16 17dtype = torch.bfloat1618device = "cuda" if torch.cuda.is_available() else "cpu"19 20MAX_SEED = np.iinfo(np.int32).max21MAX_IMAGE_SIZE = 102422 23hf_client = InferenceClient(24 api_key=os.environ.get("HF_TOKEN"),25)26VLM_MODEL = "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"27 28SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.29 30Guidelines:311. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.322. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.333. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.34 35Output only the revised prompt and nothing else."""36 37SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.2 by Black Forest Labs, an image-editing expert. You convert editing requests into one concise instruction (50-80 words, ~30 for brief requests).38 39Rules:40- Single instruction only, no commentary41- Use clear, analytical language (avoid "whimsical," "cascading," etc.)42- Specify what changes AND what stays the same (face, lighting, composition)43- Reference actual image elements44- Turn negatives into positives ("don't change X" → "keep X")45- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")46- Keep content PG-1347 48Output only the final instruction in plain text and nothing else."""49 50# Model repository IDs for 9B51REPO_ID_REGULAR = "black-forest-labs/FLUX.2-klein-base-9B"52REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-9B"53 54# Load both 9B models55print("Loading 9B Regular model...")56pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)57pipe_regular.to("cuda")58 59print("Loading 9B Distilled model...")60pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)61pipe_distilled.to("cuda")62 63# Dictionary for easy access64pipes = {65 "Distilled (4 steps)": pipe_distilled,66 "Base (50 steps)": pipe_regular,67}68 69# Default steps for each mode70DEFAULT_STEPS = {71 "Distilled (4 steps)": 4,72 "Base (50 steps)": 50,73}74 75DEFAULT_CFG = {76 "Distilled (4 steps)": 1.0,77 "Base (50 steps)": 4.0,78}79 80def image_to_data_uri(img):81 """82 Convert a PIL Image to a base64 data URI.83 84 Args:85 img: The PIL Image to convert.86 87 Returns:88 str: A data URI string containing the base64-encoded PNG image.89 """90 buffered = io.BytesIO()91 img.save(buffered, format="PNG")92 img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")93 return f"data:image/png;base64,{img_str}"94 95 96def upsample_prompt_logic(prompt, image_list):97 """98 Enhance a text prompt using a Vision-Language Model.99 100 Args:101 prompt (str): The original text prompt to enhance.102 image_list: Optional list of PIL Images for context-aware enhancement.103 104 Returns:105 str: The enhanced prompt, or the original prompt if enhancement fails.106 """107 try:108 if image_list and len(image_list) > 0:109 # Image + Text Editing Mode110 system_content = SYSTEM_PROMPT_WITH_IMAGES111 112 # Construct user message with text and images113 user_content = [{"type": "text", "text": prompt}]114 115 for img in image_list:116 data_uri = image_to_data_uri(img)117 user_content.append({118 "type": "image_url",119 "image_url": {"url": data_uri}120 })121 122 messages = [123 {"role": "system", "content": system_content},124 {"role": "user", "content": user_content}125 ]126 else:127 # Text Only Mode128 system_content = SYSTEM_PROMPT_TEXT_ONLY129 messages = [130 {"role": "system", "content": system_content},131 {"role": "user", "content": prompt}132 ]133 134 completion = hf_client.chat.completions.create(135 model=VLM_MODEL,136 messages=messages,137 max_tokens=1024138 )139 140 return completion.choices[0].message.content141 except Exception as e:142 print(f"Upsampling failed: {e}")143 return prompt144 145 146def update_dimensions_from_image(image_list):147 """148 Update width/height based on uploaded image aspect ratio.149 150 Keeps one side at 1024 and scales the other proportionally,151 with both sides as multiples of 8.152 153 Args:154 image_list: Gallery list of tuples (image, caption) from Gradio.155 156 Returns:157 tuple: A tuple of (width, height) integers, both multiples of 8.158 """159 if image_list is None or len(image_list) == 0:160 return 1024, 1024 # Default dimensions161 162 # Get the first image to determine dimensions163 img = image_list[0][0] # Gallery returns list of tuples (image, caption)164 img_width, img_height = img.size165 166 aspect_ratio = img_width / img_height167 168 if aspect_ratio >= 1: # Landscape or square169 new_width = 1024170 new_height = int(1024 / aspect_ratio)171 else: # Portrait172 new_height = 1024173 new_width = int(1024 * aspect_ratio)174 175 # Round to nearest multiple of 8176 new_width = round(new_width / 8) * 8177 new_height = round(new_height / 8) * 8178 179 # Ensure within valid range (minimum 256, maximum 1024)180 new_width = max(256, min(1024, new_width))181 new_height = max(256, min(1024, new_height))182 183 return new_width, new_height184 185 186def update_steps_from_mode(mode_choice):187 """188 Update inference steps and guidance scale based on the selected mode.189 190 Args:191 mode_choice (str): The selected mode, either "Distilled (4 steps)" or "Base (50 steps)".192 193 Returns:194 tuple: A tuple of (num_inference_steps, guidance_scale).195 """196 return DEFAULT_STEPS[mode_choice], DEFAULT_CFG[mode_choice]197 198 199@spaces.GPU(duration=85)200def infer(201 prompt: str,202 input_images=None,203 mode_choice: str = "Distilled (4 steps)",204 seed: int = 42,205 randomize_seed: bool = False,206 width: int = 1024,207 height: int = 1024,208 num_inference_steps: int = 4,209 guidance_scale: float = 4.0,210 prompt_upsampling: bool = False,211 progress=gr.Progress(track_tqdm=True)212):213 """214 Generate or edit images using FLUX.2 Klein 9B model.215 216 This tool can generate images from text prompts, or edit/combine existing images217 based on text instructions. Use the distilled mode for fast 4-step generation,218 or base mode for higher quality 50-step generation.219 220 Args:221 prompt (str): Text description of the image to generate, or editing instructions when input images are provided.222 input_images: Optional list of input images for editing or combining. Provide image URLs.223 mode_choice (str): Model mode - "Distilled (4 steps)" for fast generation or "Base (50 steps)" for higher quality.224 seed (str): Random seed for reproducible generation. Use "0" with randomize_seed=True for random results.225 randomize_seed (str): Set to "true" to use a random seed, "false" to use the specified seed.226 width (str): Output image width in pixels (256-1024, must be multiple of 8).227 height (str): Output image height in pixels (256-1024, must be multiple of 8).228 num_inference_steps (str): Number of denoising steps. Use "4" for distilled mode, "50" for base mode.229 guidance_scale (str): How closely to follow the prompt. Use "1.0" for distilled, "4.0" for base mode.230 prompt_upsampling (str): Set to "true" to automatically enhance the prompt using a VLM.231 232 Returns:233 tuple: A tuple containing the generated PIL Image and the seed used.234 """235 # Convert string inputs to proper types for MCP compatibility236 if isinstance(seed, str):237 seed = int(seed)238 if isinstance(randomize_seed, str):239 randomize_seed = randomize_seed.lower() == "true"240 if isinstance(width, str):241 width = int(width)242 if isinstance(height, str):243 height = int(height)244 if isinstance(num_inference_steps, str):245 num_inference_steps = int(num_inference_steps)246 if isinstance(guidance_scale, str):247 guidance_scale = float(guidance_scale)248 if isinstance(prompt_upsampling, str):249 prompt_upsampling = prompt_upsampling.lower() == "true"250 251 if randomize_seed:252 seed = random.randint(0, MAX_SEED)253 254 # Select the appropriate pipeline based on mode choice255 pipe = pipes[mode_choice]256 257 # Prepare image list (convert None or empty gallery to None)258 image_list = None259 if input_images is not None and len(input_images) > 0:260 image_list = []261 for item in input_images:262 image_list.append(item[0])263 264 # 1. Upsampling (Network bound)265 final_prompt = prompt266 if prompt_upsampling:267 progress(0.1, desc="Upsampling prompt...")268 final_prompt = upsample_prompt_logic(prompt, image_list)269 print(f"Original Prompt: {prompt}")270 print(f"Upsampled Prompt: {final_prompt}")271 272 # 2. Image Generation273 progress(0.2, desc=f"Generating image with 9B {mode_choice}...")274 275 generator = torch.Generator(device=device).manual_seed(seed)276 277 pipe_kwargs = {278 "prompt": final_prompt,279 "height": height,280 "width": width,281 "num_inference_steps": num_inference_steps,282 "guidance_scale": guidance_scale,283 "generator": generator,284 }285 286 # Add images if provided287 if image_list is not None:288 pipe_kwargs["image"] = image_list289 290 image = pipe(**pipe_kwargs).images[0]291 292 return image, seed293 294 295examples = [296 ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],297 ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],298 ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],299 ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],300]301 302examples_images = [303 ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]304]305 306css = """307#col-container {308 margin: 0 auto;309 max-width: 1200px;310}311.gallery-container img{312 object-fit: contain;313}314"""315 316with gr.Blocks(css=css) as demo:317 318 with gr.Column(elem_id="col-container"):319 gr.Markdown(f"""# FLUX.2 [Klein] - 9B320FLUX.2 [Klein] is a distilled model capable of generating, editing and combining images based on text instructions [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B)], [[blog](https://bfl.ai/blog/flux-2)]321 """)322 with gr.Row():323 with gr.Column():324 with gr.Row():325 prompt = gr.Text(326 label="Prompt",327 show_label=False,328 max_lines=2,329 placeholder="Enter your prompt",330 container=False,331 scale=3332 )333 334 run_button = gr.Button("Run", scale=1)335 336 with gr.Accordion("Input image(s) (optional)", open=False):337 input_images = gr.Gallery(338 label="Input Image(s)",339 type="pil",340 columns=3,341 rows=1,342 )343 344 mode_choice = gr.Radio(345 label="Mode",346 choices=["Distilled (4 steps)", "Base (50 steps)"],347 value="Distilled (4 steps)",348 )349 350 with gr.Accordion("Advanced Settings", open=False):351 352 prompt_upsampling = gr.Checkbox(353 label="Prompt Upsampling",354 value=False,355 info="Automatically enhance the prompt using a VLM"356 )357 358 seed = gr.Slider(359 label="Seed",360 minimum=0,361 maximum=MAX_SEED,362 step=1,363 value=0,364 )365 366 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)367 368 with gr.Row():369 370 width = gr.Slider(371 label="Width",372 minimum=256,373 maximum=MAX_IMAGE_SIZE,374 step=8,375 value=1024,376 )377 378 height = gr.Slider(379 label="Height",380 minimum=256,381 maximum=MAX_IMAGE_SIZE,382 step=8,383 value=1024,384 )385 386 with gr.Row():387 388 num_inference_steps = gr.Slider(389 label="Number of inference steps",390 minimum=1,391 maximum=100,392 step=1,393 value=4,394 )395 396 guidance_scale = gr.Slider(397 label="Guidance scale",398 minimum=0.0,399 maximum=10.0,400 step=0.1,401 value=1.0,402 )403 404 405 with gr.Column():406 result = gr.Image(label="Result", show_label=False)407 408 409 gr.Examples(410 examples=examples,411 fn=infer,412 inputs=[prompt],413 outputs=[result, seed],414 cache_examples=True,415 cache_mode="lazy"416 )417 418 gr.Examples(419 examples=examples_images,420 fn=infer,421 inputs=[prompt, input_images],422 outputs=[result, seed],423 cache_examples=True,424 cache_mode="lazy"425 )426 427 # Auto-update dimensions when images are uploaded428 input_images.upload(429 fn=update_dimensions_from_image,430 inputs=[input_images],431 outputs=[width, height]432 )433 434 # Auto-update steps when mode changes435 mode_choice.change(436 fn=update_steps_from_mode,437 inputs=[mode_choice],438 outputs=[num_inference_steps, guidance_scale]439 )440 441 gr.on(442 triggers=[run_button.click, prompt.submit],443 fn=infer,444 inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],445 outputs=[result, seed],446 api_name="generate" # Explicit API name for MCP tool447 )448 449# Launch with MCP server enabled450demo.launch(mcp_server=True)