from __future__ import annotations import os # Long image-conditioned prefixes cause large transient allocations; expandable # segments keep the caching allocator from fragmenting into a hard failure # (NVML_SUCCESS == r INTERNAL ASSERT FAILED in CUDACachingAllocator). # MUST be set before torch is imported. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # MUST be imported before torch / transformers / sensenova_u1 import random import threading import time import gradio as gr import numpy as np import torch from PIL import Image from transformers import AutoConfig, AutoModel, AutoTokenizer import sensenova_u1 from sensenova_u1.models.neo_unify.utils import smart_resize MODEL_ID = "sensenova/SenseNova-U1.5-8B-MoT" NORM_MEAN = (0.5, 0.5, 0.5) NORM_STD = (0.5, 0.5, 0.5) # U1.5 trained T2I aspect-ratio buckets (from the upstream examples/t2i/inference.py # SUPPORTED_RESOLUTIONS table). T2I_RESOLUTIONS: dict[str, tuple[int, int]] = { "1:1": (2048, 2048), "16:9": (2720, 1536), "9:16": (1536, 2720), "3:2": (2496, 1664), "2:3": (1664, 2496), "4:3": (2368, 1760), "3:4": (1760, 2368), "1:2": (1440, 2880), "2:1": (2880, 1440), "1:3": (1152, 3456), "3:1": (3456, 1152), } # Reference config for SenseNova-U1.5 (from the model card Quick Start): # cfg_scale=4.0, timestep_shift=3.0, num_steps=50 DEFAULT_CFG_SCALE = 4.0 DEFAULT_TIMESTEP_SHIFT = 3.0 DEFAULT_IMG_CFG_SCALE = 1.0 REFERENCE_NUM_STEPS = 50 # A fixed-seed A/B against the 50-step reference showed 28 steps keeps # composition, prompt adherence and text rendering intact — the only cost is # some micro-texture in landscape/skin — while running ~1.8x faster. Editing is # near step-invariant (0.97 SSIM vs 50 steps). Below ~24 steps structure starts # to smear, so the slider floor stays well above that. DEFAULT_NUM_STEPS = 28 MIN_NUM_STEPS = 8 MAX_NUM_STEPS = REFERENCE_NUM_STEPS # --- NCII prompt guard ------------------------------------------------------ # hfmlsoc/ncii-guard-v02 flags image-editing prompts that ask to strip, undress # or otherwise sexualize a person in a photograph. It runs on CPU in a separate # subprocess (see ncii_guard.py), and is consulted before any GPU worker is # allocated, so a refused prompt costs no GPU time. # # Threshold: the model deliberately ships no default. On its own held-out sweep # (eval/threshold-sweep.json, 980 prompts / 70 positives) precision climbs from # 0.870 at 0.5 to 0.923 at 0.80 while recall stays flat at 0.857 — so anything # below ~0.7 just donates false positives. Scores on real prompts are strongly # bimodal (0.000 or 1.000), so the threshold only decides the obfuscated # mid-band: character-separated evasions ("r e m o v e h e r d r e s s") land # at 0.73. 0.70 catches those for the cost of one extra false positive per 980 # clean prompts vs 0.80. GUARD_ID = "hfmlsoc/ncii-guard-v02" GUARD_THRESHOLD = 0.70 # The classifier is trained on *edit* prompts, and the NCII harm needs a real # photo of a real person, so only the editing path is screened. Screening plain # text-to-image costs real false positives (e.g. "classical marble statue of a # nude figure in a museum" scores 0.9999) for no gain in NCII coverage. GUARD_SCREEN_TEXT_TO_IMAGE = False # The classifier runs in a subprocess (see ncii_guard.py): loaded in this # process it leaves CUDA driver state behind that kills every subsequent # ZeroGPU worker at worker_init with "No CUDA GPUs are available". # Editing output grid factor (= patch_size * merge_size = 32). EDIT_GRID_FACTOR = 32 EDIT_TARGET_PIXELS = 2048 * 2048 # Total input-image pixel budget for the editing prefix, shared across all # reference images. The it2i prefix forward uses eager attention, whose memory # grows quadratically with the number of input image tokens; capping the total # keeps a multi-image edit from blowing past the ZeroGPU slot. EDIT_INPUT_TOTAL_MAX_PIXELS = 2048 * 2048 EDIT_INPUT_MIN_PIXELS = 512 * 512 MAX_INPUT_IMAGES = 4 MAX_SEED = 2**31 - 1 def _denorm(x: torch.Tensor) -> torch.Tensor: mean = torch.tensor(NORM_MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) std = torch.tensor(NORM_STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) return (x * std + mean).clamp(0, 1) def _to_pil(batch: torch.Tensor) -> list[Image.Image]: arr = _denorm(batch.float()).permute(0, 2, 3, 1).cpu().numpy() arr = (arr * 255.0).round().astype(np.uint8) return [Image.fromarray(a) for a in arr] def _coerce_pil(img) -> Image.Image | None: """Best-effort conversion of one Gradio gallery entry into a PIL image. Returns ``None`` for anything that isn't actually an image (stray strings such as example/placeholder values, ``None`` slots, unreadable paths...) so callers can drop it instead of crashing inside the GPU worker. """ if img is None: return None # gr.Gallery hands over (path_or_image, caption) tuples. if isinstance(img, (tuple, list)): if not img: return None img = img[0] if isinstance(img, (tuple, list)): # nested, give up return None if isinstance(img, dict): img = img.get("image") or img.get("path") or img.get("name") or img.get("url") if img is None: return None if isinstance(img, Image.Image): return img if isinstance(img, np.ndarray): try: return Image.fromarray(img.astype(np.uint8)) except Exception: return None if isinstance(img, os.PathLike): img = os.fspath(img) if isinstance(img, str): # Only treat it as an image if it really is a readable image file. if not img or not os.path.isfile(img): return None try: loaded = Image.open(img) loaded.load() return loaded except Exception: return None return None def _normalize_images(images) -> list[Image.Image]: """Turn whatever Gradio handed us into a clean list of PIL images. Runs on CPU, before the ZeroGPU worker is entered, so that a bad input is rejected up front instead of aborting a GPU call mid-flight. """ if images is None: return [] # A bare string / PIL image / array (not a list) is a single input, never # something to iterate over character by character. if isinstance(images, (str, bytes, os.PathLike, Image.Image, np.ndarray, dict)): items = [images] elif isinstance(images, (list, tuple)): items = list(images) else: items = [images] pils: list[Image.Image] = [] dropped = 0 for item in items: pil = _coerce_pil(item) if pil is None: dropped += 1 continue pils.append(pil) if dropped and not pils: raise gr.Error( "The uploaded input could not be read as an image. " "Please upload an image file, or clear the gallery for text-to-image." ) if dropped: print(f"[generate] ignored {dropped} non-image input(s).") if len(pils) > MAX_INPUT_IMAGES: raise gr.Error(f"Please use at most {MAX_INPUT_IMAGES} input images.") return pils def _input_pixel_budget(num_images: int) -> int: """Per-image pixel budget so all inputs together fit the prefix budget.""" return max(EDIT_INPUT_MIN_PIXELS, EDIT_INPUT_TOTAL_MAX_PIXELS // max(1, num_images)) def _prep_input_image(img: Image.Image, max_pixels: int) -> Image.Image: if img.mode == "RGBA": bg = Image.new("RGB", img.size, (255, 255, 255)) bg.paste(img, mask=img.split()[3]) img = bg img = img.convert("RGB") h, w = smart_resize( height=img.height, width=img.width, factor=EDIT_GRID_FACTOR, min_pixels=max_pixels, max_pixels=max_pixels, ) if (w, h) != img.size: img = img.resize((w, h), Image.LANCZOS) return img def _editing_output_size(input_img: Image.Image, target_pixels: int) -> tuple[int, int]: h, w = smart_resize( height=input_img.height, width=input_img.width, factor=EDIT_GRID_FACTOR, min_pixels=target_pixels, max_pixels=target_pixels, ) return w, h # Start the guard subprocess before the main model is loaded: Popen forks this # process, and forking it after 35GB of weights are resident is far costlier. print("[startup] loading SenseNova-U1.5-8B-MoT (this may take a few minutes)...") sensenova_u1.set_attn_backend("auto") print(f"[startup] attn backend: {sensenova_u1.effective_attn_backend()!r}") config = AutoConfig.from_pretrained(MODEL_ID) sensenova_u1.check_checkpoint_compatibility(config) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModel.from_pretrained(MODEL_ID, config=config, dtype=torch.bfloat16).to("cuda").eval() print("[startup] model ready.") import sn_aoti # Swap the 42 Qwen3DecoderLayer gen paths for one AOTI-compiled package. # No-op unless SN_AOTI=1 AND the published key matches this card exactly # (torch version, sm arch, attention backend, layer count, hidden size); any # mismatch prints one line and leaves the model eager. Nothing here touches the # GPU: the download is CPU work and the .pt2 is not opened until the first # denoising step inside the GPU worker. print(f"[startup] {sn_aoti.status()}") sn_aoti.maybe_load(model) def _estimate_duration( pil_inputs, prompt, aspect_ratio, seed, num_steps=DEFAULT_NUM_STEPS, *args, **kwargs ): # Sampling cost is ~linear in step count, so scale the 50-step baselines # instead of always reserving the worst case. try: steps = max(MIN_NUM_STEPS, min(int(num_steps), MAX_NUM_STEPS)) except (TypeError, ValueError): steps = DEFAULT_NUM_STEPS # Editing is heavier (image conditioning); give it more headroom. if pil_inputs: base = 180 else: # T2I at 2048x2048 with 50 steps is the heaviest t2i case w, h = T2I_RESOLUTIONS.get(aspect_ratio, (2048, 2048)) base = 180 if w * h > 2048 * 2048 else 120 # The prefix pass and pixel decode don't shrink with the step count, so keep # a fixed floor and scale only the sampling part. return int(30 + (base - 30) * steps / REFERENCE_NUM_STEPS) def _register_step_hook(counter: dict) -> object | None: """Count denoising steps by hooking the per-step timestep embedder. The sampling loops in ``t2i_generate`` / ``it2i_generate`` call ``fm_modules['timestep_embedder']`` exactly once per step, so a forward hook on it is a cheap, non-invasive progress signal. """ try: embedder = model.fm_modules["timestep_embedder"] except Exception: # pragma: no cover - defensive return None def _tick(_module, _args, _output): counter["step"] += 1 try: return embedder.register_forward_hook(_tick) except Exception: # pragma: no cover - defensive return None @spaces.GPU(duration=_estimate_duration) def _generate_on_gpu( pil_inputs: list, prompt: str, aspect_ratio: str, seed: int, num_steps: int = DEFAULT_NUM_STEPS, cfg_scale: float = DEFAULT_CFG_SCALE, timestep_shift: float = DEFAULT_TIMESTEP_SHIFT, img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE, ): """GPU worker. Inputs are already validated / preprocessed on CPU. Runs sampling on a background thread inside the GPU fork and yields ``("progress", steps_done, total_steps)`` ticks while it works, then a final ``("result", image)``. The caller turns those ticks into a ``gr.Progress`` bar on the output component. """ torch.cuda.reset_peak_memory_stats() counter = {"step": 0} total_steps = int(num_steps) hook = _register_step_hook(counter) result: dict = {} def _run(): try: with torch.inference_mode(): if not pil_inputs: width, height = T2I_RESOLUTIONS[aspect_ratio] tensor = model.t2i_generate( tokenizer, prompt, image_size=(width, height), cfg_scale=float(cfg_scale), cfg_norm="none", timestep_shift=float(timestep_shift), cfg_interval=(0.0, 1.0), num_steps=total_steps, batch_size=1, seed=int(seed), think_mode=False, ) else: out_w, out_h = _editing_output_size(pil_inputs[0], EDIT_TARGET_PIXELS) tensor = model.it2i_generate( tokenizer, prompt, pil_inputs, image_size=(out_w, out_h), cfg_scale=float(cfg_scale), img_cfg_scale=float(img_cfg_scale), cfg_norm="none", timestep_shift=float(timestep_shift), cfg_interval=(0.0, 1.0), num_steps=total_steps, batch_size=1, think_mode=False, seed=int(seed), ) result["image"] = _to_pil(tensor)[0] del tensor except BaseException as exc: # re-raised on the generator thread below result["error"] = exc worker = threading.Thread(target=_run, daemon=True) started = time.time() try: worker.start() while True: worker.join(0.4) done = min(counter["step"], total_steps) if not worker.is_alive(): break yield ("progress", done, total_steps, time.time() - started) print(f"[generate] peak GPU memory: {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB") exc = result.get("error") if isinstance(exc, torch.cuda.OutOfMemoryError): # pragma: no cover raise gr.Error( "Ran out of GPU memory for this request. Try a smaller input image " "or a smaller aspect ratio." ) from exc if exc is not None: raise exc yield ("result", result["image"], total_steps, time.time() - started) finally: if hook is not None: hook.remove() # Never leave a failed call's tensors behind for the next request. torch.cuda.empty_cache() def generate( images: list | None, prompt: str, aspect_ratio: str = "1:1", seed: int = 42, randomize_seed: bool = True, num_steps: int = DEFAULT_NUM_STEPS, cfg_scale: float = DEFAULT_CFG_SCALE, timestep_shift: float = DEFAULT_TIMESTEP_SHIFT, img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE, progress=gr.Progress(), ): """Generate an image from a text prompt, or edit an uploaded image. Args: images: optional uploaded image(s) to edit; leave empty for text-to-image. prompt: what to generate, or the edit instruction to apply to the input image. aspect_ratio: output aspect ratio for text-to-image (ignored when editing). seed: RNG seed for reproducible sampling. randomize_seed: if True, pick a fresh random seed each run. num_steps: denoising steps; 50 is the model card reference, 28 is ~1.8x faster. cfg_scale: prompt guidance strength (reference: 4.0). timestep_shift: flow-matching schedule shift (reference: 3.0). img_cfg_scale: input-image guidance strength, editing only (reference: 1.0). """ # Everything below runs on CPU: reject bad inputs *before* a GPU worker is # allocated, so a malformed request can never abort a GPU call mid-flight. if not isinstance(prompt, str) or not prompt.strip(): raise gr.Error("Please enter a prompt.") prompt = prompt.strip() if aspect_ratio not in T2I_RESOLUTIONS: aspect_ratio = "1:1" try: seed = int(seed) except (TypeError, ValueError): seed = 42 if randomize_seed: seed = random.randint(0, MAX_SEED) seed = max(0, min(int(seed), MAX_SEED)) def _clamp(value, low, high, fallback): try: value = float(value) except (TypeError, ValueError): return fallback if value != value: # NaN return fallback return max(low, min(value, high)) num_steps = int(_clamp(num_steps, MIN_NUM_STEPS, MAX_NUM_STEPS, DEFAULT_NUM_STEPS)) cfg_scale = _clamp(cfg_scale, 1.0, 10.0, DEFAULT_CFG_SCALE) timestep_shift = _clamp(timestep_shift, 0.5, 6.0, DEFAULT_TIMESTEP_SHIFT) img_cfg_scale = _clamp(img_cfg_scale, 1.0, 4.0, DEFAULT_IMG_CFG_SCALE) pil_images = _normalize_images(images) # Screen on CPU, before any GPU worker is allocated budget = _input_pixel_budget(len(pil_images)) pil_inputs = [_prep_input_image(img, budget) for img in pil_images] if pil_inputs: sizes = ", ".join(f"{im.width}x{im.height}" for im in pil_inputs) print(f"[generate] editing with {len(pil_inputs)} input image(s) at {sizes}") # Drive a progress bar on the output component from the GPU worker's ticks. estimated = _estimate_duration(pil_inputs, prompt, aspect_ratio, seed, num_steps) progress(0.0, desc="Starting generation…") image_out = None for kind, payload, total_steps, elapsed in _generate_on_gpu( pil_inputs, prompt, aspect_ratio, seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale, ): if kind == "progress": steps_done = int(payload) if steps_done > 0: progress( min(steps_done / max(total_steps, 1), 1.0), desc=f"Denoising step {steps_done}/{total_steps}", ) else: # Prefix / conditioning pass, before the first denoising step. progress( min(elapsed / max(estimated, 1), 0.05), desc="Encoding prompt…", ) else: image_out = payload progress(1.0, desc="Done") if image_out is None: # pragma: no cover - worker always yields a result raise gr.Error("Generation produced no image. Please try again.") return image_out, seed def generate_t2i( prompt: str, aspect_ratio: str = "1:1", seed: int = 42, randomize_seed: bool = True, num_steps: int = DEFAULT_NUM_STEPS, cfg_scale: float = DEFAULT_CFG_SCALE, timestep_shift: float = DEFAULT_TIMESTEP_SHIFT, progress=gr.Progress(), ): """Text-to-image only: generate an image from a prompt (no input image). Args: prompt: what to generate. aspect_ratio: output aspect ratio. seed: RNG seed for reproducible sampling. randomize_seed: if True, pick a fresh random seed each run. num_steps: denoising steps (default 28; 50 is the reference config). cfg_scale: prompt guidance strength. timestep_shift: flow-matching schedule shift. """ # Wired to the text-to-image gr.Examples: those rows carry (prompt, # aspect_ratio) only, so they must NOT be bound straight to `generate`, # whose first argument is the input-image list. return generate( None, prompt, aspect_ratio, seed, randomize_seed, num_steps, cfg_scale, timestep_shift, progress=progress, ) def generate_edit( images: list | None, prompt: str, seed: int = 42, randomize_seed: bool = True, num_steps: int = DEFAULT_NUM_STEPS, cfg_scale: float = DEFAULT_CFG_SCALE, timestep_shift: float = DEFAULT_TIMESTEP_SHIFT, img_cfg_scale: float = DEFAULT_IMG_CFG_SCALE, progress=gr.Progress(), ): """Image editing: apply an edit instruction to the uploaded image(s). Args: images: input image(s) to edit. prompt: the edit instruction. seed: RNG seed for reproducible sampling. randomize_seed: if True, pick a fresh random seed each run. num_steps: denoising steps (default 28; 50 is the reference config). cfg_scale: prompt guidance strength. timestep_shift: flow-matching schedule shift. img_cfg_scale: input-image guidance strength. """ return generate( images, prompt, "1:1", seed, randomize_seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale, progress=progress, ) # T2I examples: prompt + aspect ratio T2I_EXAMPLES = [ [ "A cinematic mountain lake at sunrise, realistic photography, golden mist over still water, snow-capped peaks reflected in the lake, ultra-detailed.", "1:1", ], [ 'A neon bar sign that clearly reads "OPEN LATE", dark interior, moody reflections, easy text rendering.', "16:9", ], [ "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look.", "2:3", ], [ "A cute fluffy corgi puppy wearing a tiny chef's hat, sitting at a wooden table with fresh-baked cookies, warm kitchen lighting, photorealistic.", "1:1", ], [ "Lavender fields stretching to the horizon under a pastel sunset, a small stone farmhouse, highly detailed flowers, romantic countryside scene.", "4:3", ], ] # Editing examples: gallery input (list of paths) + prompt EDIT_EXAMPLES = [ [["examples/edit_1.webp"], "Change the jacket of the person on the left to bright yellow."], [["examples/edit_2.webp"], "Make the person in the image smile."], [["examples/edit_3.webp"], "Add a bouquet of flowers."], [["examples/edit_4.webp"], "Turn the image into an American comic style."], ] CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(title="SenseNova-U1.5-8B-MoT") as demo: gr.Markdown( """ # SenseNova-U1.5-8B-MoT Unified text-to-image **and** image editing with [**SenseNova-U1.5-8B-MoT**](https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT), a natively unified multimodal model built on the [NEO-unify](https://huggingface.co/blog/sensenova/neo-unify) architecture. Leave the image upload empty for text-to-image, or upload an image and write an edit instruction. """ ) with gr.Row(): with gr.Column(scale=1): image_input_gallery = gr.Gallery( label="Upload image(s) to edit (leave empty for text-to-image)", file_types=["image"], columns=4, ) prompt_input = gr.Textbox( label="Prompt", placeholder="Describe the image to generate, or how to edit your input.", lines=3, ) aspect_ratio = gr.Dropdown( label="Aspect ratio (text-to-image only — editing keeps input ratio)", choices=list(T2I_RESOLUTIONS.keys()), value="1:1", ) generate_button = gr.Button("Generate", variant="primary") with gr.Accordion("Advanced options", open=False): num_steps = gr.Slider( label="Denoising steps", info=( f"{DEFAULT_NUM_STEPS} keeps composition and prompt adherence while " f"running ~1.8x faster than the {REFERENCE_NUM_STEPS}-step reference; " "raise it for fine texture in landscapes and skin." ), minimum=MIN_NUM_STEPS, maximum=MAX_NUM_STEPS, step=1, value=DEFAULT_NUM_STEPS, ) cfg_scale = gr.Slider( label="Guidance scale (CFG)", info="How strictly to follow the prompt. Reference: 4.0.", minimum=1.0, maximum=10.0, step=0.1, value=DEFAULT_CFG_SCALE, ) timestep_shift = gr.Slider( label="Timestep shift", info=( "Shifts sampling toward high-noise steps. Reference: 3.0 — lowering it " "changes the image rather than sharpening it." ), minimum=0.5, maximum=6.0, step=0.1, value=DEFAULT_TIMESTEP_SHIFT, ) img_cfg_scale = gr.Slider( label="Image guidance (editing only)", info="How closely an edit sticks to the input image. Reference: 1.0.", minimum=1.0, maximum=4.0, step=0.1, value=DEFAULT_IMG_CFG_SCALE, ) with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42 ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Column(scale=1): output_image = gr.Image( label="Output", type="pil", format="png", interactive=False, ) used_seed = gr.Number(label="Seed used", interactive=False) with gr.Accordion("Image Editing Examples", open=True): gr.Examples( examples=EDIT_EXAMPLES, inputs=[image_input_gallery, prompt_input], outputs=[output_image, used_seed], fn=generate_edit, cache_examples=False, run_on_click=True, ) with gr.Accordion("Text-to-Image Examples", open=True): gr.Examples( examples=T2I_EXAMPLES, inputs=[prompt_input, aspect_ratio], outputs=[output_image, used_seed], fn=generate_t2i, cache_examples=False, run_on_click=True, ) generate_button.click( fn=generate, inputs=[ image_input_gallery, prompt_input, aspect_ratio, seed, randomize_seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale, ], outputs=[output_image, used_seed], api_name="generate", ) prompt_input.submit( fn=generate, inputs=[ image_input_gallery, prompt_input, aspect_ratio, seed, randomize_seed, num_steps, cfg_scale, timestep_shift, img_cfg_scale, ], outputs=[output_image, used_seed], api_name="generate_submit", ) if __name__ == "__main__": demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)