import argparse import base64 import io import time import torch import uvicorn import gc import asyncio import traceback from typing import List, Optional, Union from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, UploadFile, File, Form from pydantic import BaseModel from PIL import Image, ImageOps # Argument parsing parser = argparse.ArgumentParser(description="Flux Image Edit Server with Nunchaku") parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") parser.add_argument("--port", type=int, default=8000, help="Port to bind to") parser.add_argument("--model", type=str, default="black-forest-labs/FLUX.1-Kontext-dev", help="Path or Repo ID of the base model") parser.add_argument("--optimized-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file") parser.add_argument("--optimized-edit-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file for editing (optional)") parser.add_argument("--backend", type=str, default="kontext", choices=["kontext", "flux2", "qwen", "glm", "zimage"], help="Backend to use: 'kontext', 'flux2', 'qwen', 'glm', or 'zimage'") parser.add_argument("--steps", type=int, default=28, help="Default number of inference steps") parser.add_argument("--guidance-scale", type=float, default=3.5, help="Default guidance scale") parser.add_argument("--qwenimage", action="store_true", help="Use QwenImageBackend (T2I only) instead of full Qwen edit backend") parser.add_argument("--uma", action="store_true", help="Enable Unified Memory Architecture mode (load all to GPU, disable offload)") parser.add_argument( "--nvfp4-text-encoder", type=str, default=None, help=( "Path to an NVFP4-pack-quantized HuggingFace text encoder " "(compressed-tensors format). Currently honoured by the zimage backend; " "swaps in vLLM's W4A4 NVFP4 CUTLASS GEMM for ~4x text-encoder VRAM savings." ), ) args = parser.parse_args() @asynccontextmanager async def lifespan(app: FastAPI): # Startup logic load_model() yield # Shutdown logic (if any) could go here app = FastAPI(lifespan=lifespan) # Global components IMAGE_DIMENSION_ALIGNMENT = 32 pipeline = None edit_pipeline = None request_lock = asyncio.Lock() is_sleeping_flag = False sleep_requested = False def load_model(): global pipeline, edit_pipeline try: if args.backend == "kontext": import KontextBackend print(f"Initializing KontextBackend...") backend = KontextBackend.KontextBackend(args.model, args.optimized_model) pipeline, edit_pipeline = backend.load() elif args.backend == "flux2": import Flux2Backend print(f"Initializing Flux2Backend...") backend = Flux2Backend.Flux2Backend(args.model) pipeline, edit_pipeline = backend.load() elif args.backend == "glm": import GlmBackend print(f"Initializing GlmBackend...") # Use provided model or default to the one in the snippet if args.model is generic # The user might pass the specific GLM model via --model, or we default in GlmBackend. # Let's pass args.model if it's not the default flux one, otherwise let GlmBackend use its default. model_to_use = args.model if args.model != "black-forest-labs/FLUX.1-Kontext-dev" else "Disty0/GLM-Image-SDNQ-4bit-dynamic" backend = GlmBackend.GlmBackend(model_to_use) pipeline, edit_pipeline = backend.load() elif args.backend.startswith("qwen"): if args.qwenimage: import QwenImageBackend print(f"Initializing QwenImageBackend (T2I only)...") backend = QwenImageBackend.QwenImageBackend(args.model, args.optimized_model) pipeline, edit_pipeline = backend.load() else: import QwenBackend print(f"Initializing QwenBackend...") backend = QwenBackend.QwenBackend(args.model, args.optimized_model, optimized_edit_model_path=args.optimized_edit_model, uma=args.uma) pipeline, edit_pipeline = backend.load() elif args.backend == "zimage": import ZImageTurboBackend print(f"Initializing ZImageTurboBackend...") backend = ZImageTurboBackend.ZImageTurboBackend( args.model, args.optimized_model, uma=args.uma, nvfp4_text_encoder_path=args.nvfp4_text_encoder, ) pipeline, edit_pipeline = backend.load() else: raise ValueError(f"Unknown backend: {args.backend}") except Exception as e: print(f"Oh no! The model refused to wake up: {e}") raise e # Enable progress bar for diffusers import diffusers.utils.logging diffusers.utils.logging.enable_progress_bar() diffusers.utils.logging.set_verbosity_info() print("Model loaded successfully! Ready for editing quests!") def flush(): gc.collect() torch.cuda.empty_cache() class ImageGenerationRequest(BaseModel): prompt: str n: int = 1 size: str = "1024x1024" response_format: str = "b64_json" quality: str = "standard" style: str = "vivid" num_inference_steps: Optional[int] = None guidance_scale: Optional[float] = None negative_prompt: Optional[str] = None seed: Optional[int] = None @app.post("/v1/sleep") async def sleep_endpoint(): global is_sleeping_flag, sleep_requested sleep_requested = True try: async with request_lock: if not is_sleeping_flag and sleep_requested: print("Sleep requested, moving models to CPU...") for p in [pipeline, edit_pipeline]: if not p: continue for name, component in p.components.items(): if isinstance(component, torch.nn.Module): # Special handling for Nunchaku which blocks .to() if offload is True if hasattr(component, "set_offload") and getattr(component, "offload", False): component.set_offload(False) component._nunchaku_was_offloaded = True try: component.to("cpu") except Exception as e: pass flush() is_sleeping_flag = True finally: sleep_requested = False return {"status": "sleep completed", "is_sleeping": is_sleeping_flag} @app.post("/v1/wake_up") async def wake_up_endpoint(): global is_sleeping_flag, sleep_requested sleep_requested = False async with request_lock: if is_sleeping_flag: print("Waking up, restoring models to CUDA...") for p in [pipeline, edit_pipeline]: if not p: continue excluded = getattr(p, "_exclude_from_cpu_offload", []) for name, component in p.components.items(): if isinstance(component, torch.nn.Module): if getattr(component, "_nunchaku_was_offloaded", False): component.set_offload(True, use_pin_memory=True, num_blocks_on_gpu=8) for attr in ["img_in", "txt_in", "txt_norm", "time_text_embed", "norm_out", "proj_out"]: if hasattr(component, attr): try: getattr(component, attr).to("cuda") except Exception: pass component._nunchaku_was_offloaded = False elif not hasattr(component, "_hf_hook") or name in excluded: try: component.to("cuda") except Exception: pass is_sleeping_flag = False return {"status": "awoken", "is_sleeping": False} @app.get("/v1/is_sleeping") async def is_sleeping_endpoint(): return {"is_sleeping": is_sleeping_flag} @app.get("/v1/memory_stats") async def memory_stats_endpoint(): """Lightweight introspection endpoint that returns PyTorch's CUDA allocator snapshot. Used to diagnose VRAM/UMA bloat without restarting the server.""" stats = {} if torch.cuda.is_available(): stats["allocated_gb"] = torch.cuda.memory_allocated() / 1e9 stats["reserved_gb"] = torch.cuda.memory_reserved() / 1e9 stats["max_allocated_gb"] = torch.cuda.max_memory_allocated() / 1e9 stats["max_reserved_gb"] = torch.cuda.max_memory_reserved() / 1e9 # Top allocations by size from the allocator snapshot (>=64 MiB) try: snap = torch.cuda.memory_snapshot() blocks = [] for seg in snap: for b in seg.get("blocks", []): if b.get("state") == "active_allocated" and b.get("size", 0) >= 64 * 1024 * 1024: blocks.append(b["size"]) blocks.sort(reverse=True) stats["large_active_blocks_gb"] = [round(s / 1e9, 3) for s in blocks[:20]] stats["large_active_blocks_total_gb"] = round(sum(blocks) / 1e9, 3) stats["large_active_blocks_count"] = len(blocks) except Exception as e: stats["snapshot_error"] = str(e) # Walk Python objects to find big tensors and group them try: import gc as _gc seen = set() big = [] for obj in _gc.get_objects(): try: if isinstance(obj, torch.Tensor) and obj.is_cuda: ptr = obj.data_ptr() if ptr in seen or ptr == 0: continue seen.add(ptr) sz = obj.element_size() * obj.numel() if sz >= 16 * 1024 * 1024: big.append((sz, tuple(obj.shape), str(obj.dtype))) except Exception: continue big.sort(reverse=True) # Group by (shape, dtype) from collections import Counter grouped = Counter((shape, dtype) for _, shape, dtype in big) stats["big_tensor_groups"] = [ {"shape": list(shape), "dtype": dtype, "count": cnt, "size_gb_each": round( (1 if shape == () else (lambda l: __import__('functools').reduce(lambda a, b: a*b, l, 1))(shape)) * ( 8 if 'int64' in dtype or 'float64' in dtype else 4 if 'int32' in dtype or 'float32' in dtype else 2 if 'bfloat16' in dtype or 'float16' in dtype else 1 ) / 1e9, 4)} for (shape, dtype), cnt in grouped.most_common(30) ] stats["big_tensor_count"] = len(big) stats["big_tensor_total_gb"] = round(sum(s for s, _, _ in big) / 1e9, 3) except Exception as e: stats["walk_error"] = str(e) return stats @app.post("/v1/images/edits") async def edit_image( image: Union[List[UploadFile], UploadFile] = File(...), prompt: str = Form(...), n: int = Form(1), size: str = Form("1024x1024"), response_format: str = Form("b64_json"), # Default to b64_json guidance_scale: Optional[float] = Form(None), num_inference_steps: Optional[int] = Form(None), negative_prompt: Optional[str] = Form(None), seed: Optional[int] = Form(None) ): # Use CLI defaults if not provided steps = num_inference_steps if num_inference_steps is not None else args.steps cfg_scale = guidance_scale if guidance_scale is not None else args.guidance_scale neg_prompt = negative_prompt if negative_prompt is not None else "" # Default empty for now, or maybe None? generator = None import random if seed is None: seed = random.randint(0, 2**32 - 1) print(f"Using seed: {seed}") generator = torch.Generator(device="cuda").manual_seed(seed) if not edit_pipeline: raise HTTPException(status_code=500, detail="Model not loaded") if sleep_requested or is_sleeping_flag: raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.") async with request_lock: print(f"Received edit request: {prompt}") # Processing the input image(s) input_files = image if isinstance(image, list) else [image] init_images = [] try: for img_file in input_files: await img_file.seek(0) contents = await img_file.read() img = Image.open(io.BytesIO(contents)).convert("RGB") init_images.append(img) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid image file: {e}") if not init_images: raise HTTPException(status_code=400, detail="No images provided") # Parse max target dimensions from requested size try: target_width, target_height = map(int, size.split("x")) except ValueError: target_width, target_height = 1024, 1024 # Calculate new dimensions preserving aspect ratio based on the first image first_image = init_images[0] orig_width, orig_height = first_image.size scale = min(target_width / orig_width, target_height / orig_height) new_width = int(orig_width * scale) new_height = int(orig_height * scale) # Ensure dimensions are aligned to 32 for compatibility (e.g. GLM-Image) width = (new_width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT height = (new_height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT # Resize input images to match the calculated target size, padding if necessary resized_images = [] for img in init_images: if img.size != (width, height): # Use ImageOps.pad to preserve aspect ratio and center in the target size # This handles cases where subsequent images might have different ARs img = ImageOps.pad(img, (width, height), method=Image.LANCZOS, color=(0, 0, 0)) resized_images.append(img) # If single image, pass as item, if multiple, pass as list # GLM pipeline has a bug where it checks len() on the input, so it must be a list if len(resized_images) > 1 or args.backend == "glm": image_input = resized_images else: image_input = resized_images[0] response_images = [] try: if args.backend.startswith("qwen"): # Qwen specific parameters # guidance_scale maps to true_cfg_scale if args.qwenimage: # QwenImageBackend is T2I only, so it doesn't take an image generated_images = edit_pipeline( prompt=prompt, height=height, width=width, num_inference_steps=steps, true_cfg_scale=cfg_scale, num_images_per_prompt=n, generator=generator, ).images else: # Full Qwen edit backend takes an image (or list of images now) generated_images = edit_pipeline( image=image_input, prompt=prompt, height=height, width=width, negative_prompt=neg_prompt, num_inference_steps=steps, true_cfg_scale=cfg_scale, num_images_per_prompt=n, generator=generator, ).images else: # Standard Flux/Kontext or GLM # GLM I2I Fix: Manually move vision encoder to GPU because get_image_features escapes hooks if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"): print("Manually moving GLM Vision Encoder to GPU...") edit_pipeline.vision_language_encoder.to("cuda") try: generated_images = edit_pipeline( image=image_input, prompt=prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=cfg_scale, num_images_per_prompt=n, generator=generator, ).images finally: if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"): print("Moving GLM Vision Encoder back to CPU...") edit_pipeline.vision_language_encoder.to("cpu") for img in generated_images: buffered = io.BytesIO() img.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") if response_format == "b64_json": response_images.append({"b64_json": img_str}) else: # If url is requested we can't really do it without storage, so we fallback or error? # For now, let's just assume simple b64_json as per request response_images.append({"b64_json": img_str}) # Fallback except Exception as e: print(f"Error during editing: {e}") print(traceback.format_exc()) raise HTTPException(status_code=500, detail=str(e)) finally: flush() return { "created": int(time.time()), "data": response_images } @app.post("/v1/images/generations") async def generate_image(request: ImageGenerationRequest): if not pipeline: raise HTTPException(status_code=500, detail="Model not loaded") if sleep_requested or is_sleeping_flag: raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.") async with request_lock: #print(f"Received generation request: {request.prompt}") # Parse size try: width, height = map(int, request.size.split("x")) except ValueError: width, height = 1024, 1024 # Ensure dimensions are aligned to 32 width = (width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT height = (height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT response_images = [] try: # Generate images (no image argument for txt2img!) steps = request.num_inference_steps if request.num_inference_steps is not None else args.steps cfg_scale = request.guidance_scale if request.guidance_scale is not None else args.guidance_scale # negative_prompt not in standard request body in original snippet, but we added it to model neg_prompt = request.negative_prompt if request.negative_prompt is not None else "" generator = None import random seed = request.seed if seed is None: seed = random.randint(0, 2**32 - 1) print(f"Using seed: {seed}") generator = torch.Generator(device="cuda").manual_seed(seed) if args.backend.startswith("qwen"): generated_images = pipeline( prompt=request.prompt, height=height, width=width, num_inference_steps=steps, true_cfg_scale=cfg_scale, num_images_per_prompt=request.n, negative_prompt=neg_prompt, generator=generator, ).images else: generated_images = pipeline( prompt=request.prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=cfg_scale, num_images_per_prompt=request.n, generator=generator, # Not passing negative_prompt here for generation unless we confirm support in standard Flux pipeline? ).images for img in generated_images: buffered = io.BytesIO() img.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") response_images.append({"b64_json": img_str}) except Exception as e: print(f"Error during generation: {e}") print(traceback.format_exc()) raise HTTPException(status_code=500, detail=str(e)) finally: flush() return { "created": int(time.time()), "data": response_images } if __name__ == "__main__": uvicorn.run(app, host=args.host, port=args.port)