Spaces:
Running on Zero
Running on Zero
| # coding=utf-8 | |
| """Inference engine for OpenHTR-Gen. | |
| Conditioning: a glyph strip is stacked on top of the real crop, the target box is | |
| zeroed in the condition image, and Flux Fill regenerates the canvas. The strip is then | |
| dropped and the edited region is cropped out. | |
| Invariants (must match the training pipeline): | |
| * canvas is size x (size + strip_height_for(size)), default 512 x 608; | |
| * the target box must be zeroed in the condition image; | |
| * call the pipeline with `masked_image_latents=` so it skips `image * (1 - mask)`, | |
| which requires diffusers 0.32.x (0.33+ rejects `image=None`); | |
| * only T5 is used; CLIP is not loaded. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| import random | |
| from contextlib import contextmanager | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| from diffusers import FluxFillPipeline, FluxTransformer2DModel | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| DEFAULT_BASE_MODEL = os.environ.get("BASE_MODEL", "black-forest-labs/FLUX.1-Fill-dev") | |
| DEFAULT_FONT_PATH = os.path.join(HERE, "Arial-Unicode-Regular.ttf") | |
| LORA_FILENAME = os.environ.get("LORA_FILENAME", "pytorch_lora_weights.safetensors") | |
| # Fallback host for the style adapter when no local copy is shipped with the Space. | |
| DEFAULT_LORA_REPO = os.environ.get("DEFAULT_LORA_REPO", "OpenHTR/OpenHTR-Gen-LoRA") | |
| # Must match the training resolution. | |
| DEFAULT_SIZE = 512 | |
| _PIPE = None | |
| # ============================================================================= | |
| # Prompt | |
| # ============================================================================= | |
| CONCAT_PROMPT_TEMPLATE = ( | |
| "The pair of images highlights some white words on a black background, " | |
| "as well as their style on a real-world scene image. " | |
| "[IMAGE1] is a template image rendering the text, with the words {words}; " | |
| "[IMAGE2] shows the text content {words} naturally and correspondingly integrated into the image." | |
| ) | |
| # CLIP is not loaded, but pooled_projections is required; training used zeros. | |
| POOLED_PROJECTION_DIM = 768 | |
| def generate_prompt_concat(words): | |
| if isinstance(words, str): | |
| words = [words] | |
| words_str = ', '.join(f"'{w}'" for w in words) | |
| return CONCAT_PROMPT_TEMPLATE.format(words=words_str) | |
| def zero_pooled_projections(batch_size, device, dtype): | |
| return torch.zeros(batch_size, POOLED_PROJECTION_DIM, device=device, dtype=dtype) | |
| # ============================================================================= | |
| # Geometry | |
| # ============================================================================= | |
| def polygon_bbox(polygon): | |
| xs = [p[0] for p in polygon] | |
| ys = [p[1] for p in polygon] | |
| return min(xs), min(ys), max(xs), max(ys) | |
| def _window_range(img_w, img_h, ux0, uy0, ux1, uy1, L): | |
| """Valid top-left range so that [sx0, sx0+L) x [sy0, sy0+L) covers the box.""" | |
| sx0_min = max(0, int(math.ceil(ux1 - L))) | |
| sx0_max = min(int(math.floor(ux0)), int(img_w - L)) | |
| sy0_min = max(0, int(math.ceil(uy1 - L))) | |
| sy0_max = min(int(math.floor(uy0)), int(img_h - L)) | |
| return sx0_min, sx0_max, sy0_min, sy0_max | |
| def compute_crop_box(img_w, img_h, target_polygon, other_polygons, | |
| target_size=512, rng=None, deterministic=False): | |
| """Fixed-size square crop window. Returns (sx0, sy0, L, target_size) or None.""" | |
| if rng is None: | |
| rng = random | |
| tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon) | |
| L_max = min(img_w, img_h) | |
| if L_max <= 1: | |
| return None | |
| tw, th = tx1 - tx0, ty1 - ty0 | |
| if max(tw, th) > L_max: | |
| return None | |
| def _pick(L, ux0, uy0, ux1, uy1): | |
| sx0_min, sx0_max, sy0_min, sy0_max = _window_range(img_w, img_h, ux0, uy0, ux1, uy1, L) | |
| if sx0_max < sx0_min or sy0_max < sy0_min: | |
| return None | |
| if deterministic: | |
| return (sx0_min + sx0_max) // 2, (sy0_min + sy0_max) // 2 | |
| return rng.randint(sx0_min, sx0_max), rng.randint(sy0_min, sy0_max) | |
| # Prefer a window that also contains a second annotation box. | |
| best_L_min, best_union = None, None | |
| for op in other_polygons: | |
| ox0, oy0, ox1, oy1 = polygon_bbox(op) | |
| ux0, uy0 = min(tx0, ox0), min(ty0, oy0) | |
| ux1, uy1 = max(tx1, ox1), max(ty1, oy1) | |
| L_cand = max(ux1 - ux0, uy1 - uy0) | |
| if L_cand > L_max: | |
| continue | |
| if best_L_min is None or L_cand < best_L_min: | |
| best_L_min = L_cand | |
| best_union = (ux0, uy0, ux1, uy1) | |
| if best_union is not None: | |
| L_min = max(best_L_min, max(tw, th), 1.0) | |
| if L_min <= L_max: | |
| L = min(int(target_size), int(L_max)) if L_min <= target_size \ | |
| else min(int(math.ceil(L_min)), int(L_max)) | |
| got = _pick(L, *best_union) | |
| if got is None: | |
| L = int(math.ceil(L_min)) | |
| got = _pick(L, *best_union) | |
| if got is not None: | |
| return got[0], got[1], L, target_size | |
| # Fallback: cover the target box only. | |
| L = min(int(target_size), int(L_max)) | |
| if max(tw, th) > L: | |
| L = min(int(math.ceil(max(tw, th))), int(L_max)) | |
| got = _pick(L, tx0, ty0, tx1, ty1) | |
| if got is None: | |
| return None | |
| return got[0], got[1], L, target_size | |
| def _pad_to_fit(real_pil, target_polygon, other_polygons, needed_L): | |
| """Replicate-pad the page when the box is larger than its short side.""" | |
| W, H = real_pil.size | |
| if needed_L <= min(W, H): | |
| return real_pil, target_polygon, other_polygons, 0 | |
| pad = int(math.ceil((needed_L - min(W, H)) / 2.0)) + 1 | |
| arr = cv2.copyMakeBorder(np.array(real_pil), pad, pad, pad, pad, cv2.BORDER_REPLICATE) | |
| shifted_t = [[px + pad, py + pad] for px, py in target_polygon] | |
| shifted_o = [[[px + pad, py + pad] for px, py in op] for op in other_polygons] | |
| return Image.fromarray(arr), shifted_t, shifted_o, pad | |
| def loosen_polygon_width(polygon, other_polygons, canvas_w, max_scale=2.0, rng=None): | |
| """Random horizontal widening of the mask (training only).""" | |
| if rng is None: | |
| rng = random | |
| x0, y0, x1, y1 = polygon_bbox(polygon) | |
| width = max(x1 - x0, 1.0) | |
| cx = (x0 + x1) / 2.0 | |
| extra = width * (rng.uniform(1.0, max_scale) - 1.0) | |
| left_ratio = rng.random() | |
| desired_left = extra * left_ratio | |
| desired_right = extra * (1.0 - left_ratio) | |
| left_limit, right_limit = 0.0, float(canvas_w) | |
| for op in other_polygons: | |
| ox0, oy0, ox1, oy1 = polygon_bbox(op) | |
| if oy1 <= y0 or oy0 >= y1: | |
| continue | |
| if ox1 <= x0: | |
| left_limit = max(left_limit, ox1) | |
| elif ox0 >= x1: | |
| right_limit = min(right_limit, ox0) | |
| left_add = min(desired_left, max(0.0, x0 - left_limit)) | |
| right_add = min(desired_right, max(0.0, right_limit - x1)) | |
| out = [] | |
| for px, py in polygon: | |
| out.append([px - left_add, py] if px <= cx else [px + right_add, py]) | |
| return out | |
| def polygon_to_mask(polygon, canvas_w, canvas_h): | |
| """255 = region to rewrite.""" | |
| m = np.zeros((int(canvas_h), int(canvas_w)), dtype=np.uint8) | |
| cv2.fillPoly(m, [np.array(polygon, dtype=np.int32)], 255) | |
| return m | |
| # ============================================================================= | |
| # Condition construction (glyph strip on top, real crop below) | |
| # ============================================================================= | |
| def strip_height_for(size): | |
| """Strip height: size // 6 rounded up to a multiple of 16 (VAE x8 + 2x2 patchify).""" | |
| base = size // 6 | |
| return int(math.ceil(base / 16.0)) * 16 | |
| def concat_canvas_size(size): | |
| return size, size + strip_height_for(size) | |
| def crop_scene_part(image, size): | |
| """Drop the glyph strip, keeping the bottom size x size scene.""" | |
| w, h = image.size | |
| return image.crop((0, h - size, size, h)) | |
| def render_glyph_strip(font, text, width, height, margin=0.90, min_font_size=8, max_font_size=None): | |
| """Render white text on black, centered and scaled to fit the strip.""" | |
| canvas = Image.new("L", (int(width), int(height)), 0) | |
| if not text or not text.strip(): | |
| return canvas | |
| g_size = 50 | |
| try: | |
| probe = font.font_variant(size=g_size) | |
| except Exception: | |
| probe = font | |
| left, top, right, bottom = probe.getbbox(text) | |
| tw = max(right - left, 1) | |
| th = max(bottom - top, 1) | |
| ratio = min(width * margin / tw, height * margin / th) | |
| size = int(g_size * ratio) | |
| if max_font_size is not None: | |
| size = min(size, int(max_font_size)) | |
| size = max(size, min_font_size) | |
| try: | |
| final_font = font.font_variant(size=size) | |
| except Exception: | |
| final_font = font | |
| ImageDraw.Draw(canvas).text( | |
| (width / 2.0, height / 2.0), text, font=final_font, fill=255, anchor="mm", | |
| ) | |
| return canvas | |
| def build_sample_concat(real_pil, text, target_polygon, other_polygons, font, | |
| size=512, loose_mask=False, mask_aug=False, | |
| rng=None, deterministic=True, allow_pad=True): | |
| """Build the condition tensors for one edit. Returns a dict or None.""" | |
| if rng is None: | |
| rng = random | |
| if allow_pad: | |
| tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon) | |
| needed = max(int(size), int(math.ceil(max(tx1 - tx0, ty1 - ty0)))) | |
| real_pil, target_polygon, other_polygons, _ = _pad_to_fit( | |
| real_pil, target_polygon, other_polygons, needed | |
| ) | |
| img_w, img_h = real_pil.size | |
| crop = compute_crop_box(img_w, img_h, target_polygon, other_polygons, | |
| target_size=size, rng=rng, deterministic=deterministic) | |
| if crop is None: | |
| return None | |
| sx0, sy0, L, target_size = crop | |
| real = real_pil.crop((sx0, sy0, sx0 + L, sy0 + L)) | |
| polygon = [[px - sx0, py - sy0] for px, py in target_polygon] | |
| others = [[[px - sx0, py - sy0] for px, py in op] for op in other_polygons] | |
| if L != target_size: | |
| scale = target_size / L | |
| real = real.resize((target_size, target_size), Image.BICUBIC) | |
| polygon = [[px * scale, py * scale] for px, py in polygon] | |
| others = [[[px * scale, py * scale] for px, py in op] for op in others] | |
| w = h = target_size | |
| if loose_mask and not deterministic: | |
| polygon = loosen_polygon_width(polygon, others, canvas_w=w, rng=rng) | |
| mask_np = polygon_to_mask(polygon, w, h) | |
| if mask_aug: | |
| r = rng.random() | |
| k = np.ones((3, 3), dtype=np.uint8) | |
| if r < 0.7: | |
| pass | |
| elif r < 0.8: | |
| mask_np = cv2.dilate(mask_np, k, iterations=1) | |
| elif r < 0.9: | |
| mask_np = cv2.erode(mask_np, k, iterations=1) | |
| elif r < 0.95: | |
| mask_np = cv2.dilate(mask_np, k, iterations=2) | |
| else: | |
| mask_np = cv2.erode(mask_np, k, iterations=2) | |
| strip_h = strip_height_for(size) | |
| render_img = render_glyph_strip(font, text, w, strip_h) | |
| combined_image_np = np.vstack(( | |
| np.repeat(np.array(render_img)[..., None], 3, axis=2), | |
| np.array(real), | |
| )) | |
| combined_mask_np = np.vstack(( | |
| np.zeros((strip_h, w), dtype=np.uint8), | |
| mask_np, | |
| )) | |
| # Zero the target box in the condition image; otherwise the answer leaks in. | |
| m = (combined_mask_np > 127).astype(np.uint8)[..., None] | |
| combined_cond_image_np = (combined_image_np * (1 - m)).astype(np.uint8) | |
| return { | |
| "combined_image": Image.fromarray(combined_image_np), | |
| "combined_cond_image": Image.fromarray(combined_cond_image_np), | |
| "combined_mask": Image.fromarray(combined_mask_np), | |
| "polygon": polygon, | |
| "strip_height": strip_h, | |
| "prompt": generate_prompt_concat([text]), | |
| "text": text, | |
| "crop_scale": target_size / L, | |
| } | |
| # ============================================================================= | |
| # 320-channel condition latents (64 image + 256 mask) | |
| # ============================================================================= | |
| def _pack_latents(x, batch_size, num_channels, height, width): | |
| """Same layout as FluxFillPipeline._pack_latents.""" | |
| x = x.view(batch_size, num_channels, height // 2, 2, width // 2, 2) | |
| x = x.permute(0, 2, 4, 1, 3, 5) | |
| return x.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4) | |
| def build_masked_image_latents(vae, cond_image, mask, dtype, vae_scale_factor=8, generator=None): | |
| """Encode the condition image and pack it together with the mask.""" | |
| if mask.ndim == 3: | |
| mask = mask.unsqueeze(1) | |
| mask = (mask > 0.5).to(dtype) | |
| cond_image = cond_image.to(device=vae.device, dtype=vae.dtype) | |
| posterior = vae.encode(cond_image).latent_dist | |
| latents = posterior.sample(generator=generator) if generator is not None else posterior.sample() | |
| latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor | |
| latents = latents.to(dtype) | |
| bsz, ch, lh, lw = latents.shape | |
| packed_image = _pack_latents(latents, bsz, ch, lh, lw) | |
| m = mask.to(device=latents.device, dtype=dtype)[:, 0] | |
| m = m.view(bsz, lh, vae_scale_factor, lw, vae_scale_factor) | |
| m = m.permute(0, 2, 4, 1, 3).reshape(bsz, vae_scale_factor ** 2, lh, lw) | |
| packed_mask = _pack_latents(m, bsz, vae_scale_factor ** 2, lh, lw) | |
| return torch.cat([packed_image, packed_mask], dim=-1) | |
| # ============================================================================= | |
| # Model loading | |
| # ============================================================================= | |
| def _tolerate_redundant_to_on_quantized_models(): | |
| """Work around accelerate calling .to() on bitsandbytes models.""" | |
| try: | |
| from transformers.modeling_utils import PreTrainedModel | |
| except Exception: # pragma: no cover | |
| yield | |
| return | |
| orig_to = PreTrainedModel.to | |
| def patched_to(self, *args, **kwargs): | |
| if getattr(self, "is_loaded_in_8bit", False) or getattr(self, "is_loaded_in_4bit", False): | |
| return self | |
| return orig_to(self, *args, **kwargs) | |
| PreTrainedModel.to = patched_to | |
| try: | |
| yield | |
| finally: | |
| PreTrainedModel.to = orig_to | |
| def _as_lora_file(path): | |
| if not path: | |
| return None | |
| if os.path.isdir(path): | |
| cand = os.path.join(path, LORA_FILENAME) | |
| return cand if os.path.isfile(cand) else None | |
| return path if os.path.isfile(path) else None | |
| def _download_lora(repo_id): | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| return hf_hub_download( | |
| repo_id, LORA_FILENAME, | |
| token=os.environ.get("HF_TOKEN") or None, | |
| ) | |
| except Exception as e: # pragma: no cover | |
| print(f"[lora] could not fetch {repo_id}/{LORA_FILENAME}: {e}") | |
| return None | |
| def resolve_lora_path(): | |
| """Locate the style adapter weights. Priority: | |
| 1) LORA_PATH (local file/dir, or a HF repo id) | |
| 2) LORA_REPO_ID (HF repo id) | |
| 3) a local lora_checkpoint_v2 folder shipped with this app | |
| 4) DEFAULT_LORA_REPO (OpenHTR/OpenHTR-Gen-LoRA unless overridden) | |
| Returns None when nothing is reachable (the demo still runs on the base model). | |
| """ | |
| env_path = os.environ.get("LORA_PATH", "").strip() | |
| if env_path: | |
| hit = _as_lora_file(env_path) or _download_lora(env_path) | |
| if hit: | |
| return hit | |
| repo_id = os.environ.get("LORA_REPO_ID", "").strip() | |
| if repo_id: | |
| hit = _download_lora(repo_id) | |
| if hit: | |
| return hit | |
| candidates = [ | |
| os.path.join(HERE, "lora_checkpoint_v2"), | |
| os.path.join(HERE, "models", "lora_checkpoint_v2"), | |
| os.path.join(os.path.dirname(HERE), "models", "lora_checkpoint_v2"), | |
| os.path.join(os.path.dirname(HERE), "lora_checkpoint_v2"), | |
| ] | |
| for cand in candidates: | |
| hit = _as_lora_file(cand) | |
| if hit: | |
| return hit | |
| if DEFAULT_LORA_REPO: | |
| return _download_lora(DEFAULT_LORA_REPO) | |
| return None | |
| def load_pipeline(base_model=DEFAULT_BASE_MODEL, lora_path=None, | |
| quantize_8bit=False, device="cuda", zero_gpu=False): | |
| """Load the base model and style adapter once per process. | |
| On ZeroGPU, call this at module scope, after importing ``spaces``. Its | |
| CUDA emulation registers the weights for GPU worker reuse. Avoid the | |
| Accelerate ``device_map`` loading route in that environment. | |
| """ | |
| global _PIPE | |
| if _PIPE is not None: | |
| return _PIPE | |
| if quantize_8bit: | |
| compute_dtype = torch.float16 | |
| elif zero_gpu: | |
| # ZeroGPU exposes CUDA through an emulation layer before GPU allocation. | |
| compute_dtype = torch.bfloat16 | |
| else: | |
| compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 | |
| transformer_kwargs = dict(torch_dtype=compute_dtype) | |
| if quantize_8bit: | |
| from diffusers import BitsAndBytesConfig as DiffusersBnbConfig | |
| transformer_kwargs["quantization_config"] = DiffusersBnbConfig(load_in_8bit=True) | |
| elif not zero_gpu and torch.device(device).type == "cuda" and os.environ.get("HW_DEVICE_MAP", "1") != "0": | |
| # Direct-to-GPU loading is a useful local option, but Accelerate's | |
| # hooks collide with DiffusionPipeline.to() in diffusers 0.32.2. | |
| # ZeroGPU must use regular loading + module-scope CUDA registration. | |
| dev = torch.device(device) | |
| if dev.index is None: | |
| dev = torch.device("cuda", torch.cuda.current_device()) | |
| transformer_kwargs["device_map"] = {"": str(dev)} | |
| print(f"[load] base={base_model} dtype={compute_dtype} 8bit={quantize_8bit}") | |
| with _tolerate_redundant_to_on_quantized_models(): | |
| try: | |
| transformer = FluxTransformer2DModel.from_pretrained( | |
| base_model, subfolder="transformer", **transformer_kwargs | |
| ) | |
| except (ValueError, TypeError) as e: | |
| if "device_map" not in transformer_kwargs: | |
| raise | |
| print(f"[load] device_map failed ({e}), falling back to plain loading") | |
| transformer_kwargs.pop("device_map") | |
| torch.cuda.empty_cache() | |
| transformer = FluxTransformer2DModel.from_pretrained( | |
| base_model, subfolder="transformer", **transformer_kwargs | |
| ) | |
| if lora_path: | |
| print(f"[load] adapter: {lora_path}") | |
| state_dict, network_alphas = FluxFillPipeline.lora_state_dict( | |
| pretrained_model_name_or_path_or_dict=lora_path, return_alphas=True | |
| ) | |
| if not all("lora" in k or "dora_scale" in k for k in state_dict.keys()): | |
| raise ValueError(f"invalid adapter weights: {lora_path}") | |
| FluxFillPipeline.load_lora_into_transformer( | |
| state_dict=state_dict, network_alphas=network_alphas, transformer=transformer, | |
| ) | |
| else: | |
| print("[load] no adapter found, running with the base model only") | |
| # CLIP is not needed: text_encoder/tokenizer are skipped, T5 (text_encoder_2) stays. | |
| pipe = FluxFillPipeline.from_pretrained( | |
| base_model, | |
| transformer=transformer, | |
| text_encoder=None, | |
| tokenizer=None, | |
| torch_dtype=compute_dtype, | |
| ) | |
| pipe.vae.enable_tiling() | |
| pipe.vae.enable_slicing() | |
| if "device_map" in transformer_kwargs or quantize_8bit: | |
| # The transformer is already dispatched (or quantized). Moving the | |
| # whole pipeline would either conflict with Accelerate hooks or | |
| # attempt an unsupported .to() on a bitsandbytes module. | |
| # Only the remaining PyTorch modules need to be moved. | |
| pipe.text_encoder_2.to(device) | |
| pipe.vae.to(device) | |
| else: | |
| # ZeroGPU intercepts this call at module scope and registers the | |
| # weights for later GPU worker allocation. It also works locally. | |
| pipe.to(device) | |
| print( | |
| f"[load] devices: transformer={pipe.transformer.device}, " | |
| f"T5={pipe.text_encoder_2.device}, VAE={pipe.vae.device}", | |
| flush=True, | |
| ) | |
| pipe.set_progress_bar_config(disable=True) | |
| pipe._hw_compute_dtype = compute_dtype | |
| _PIPE = pipe | |
| return _PIPE | |
| def pipeline_dtype(pipe): | |
| return getattr(pipe, "_hw_compute_dtype", torch.bfloat16) | |
| def pipeline_device(pipe): | |
| try: | |
| return pipe._execution_device | |
| except Exception: | |
| return pipe.transformer.device | |
| # ============================================================================= | |
| # Inference | |
| # ============================================================================= | |
| def _encode_t5(pipe, prompt, max_sequence_length=512): | |
| dtype = pipeline_dtype(pipe) | |
| inputs = pipe.tokenizer_2( | |
| [prompt] if isinstance(prompt, str) else prompt, | |
| padding="max_length", max_length=max_sequence_length, truncation=True, | |
| return_length=False, return_overflowing_tokens=False, return_tensors="pt", | |
| ) | |
| device = pipe.text_encoder_2.device | |
| with torch.no_grad(): | |
| return pipe.text_encoder_2(inputs.input_ids.to(device))[0].to(dtype) | |
| def run_one_inference(pipe, cond_image, mask, prompt, num_steps, guidance_scale, seed, | |
| max_sequence_length=512): | |
| """Run Flux Fill on one condition canvas; returns a PIL image of the same size.""" | |
| assert cond_image.size == mask.size, f"cond_image{cond_image.size} vs mask{mask.size}" | |
| width, height = cond_image.size | |
| dtype = pipeline_dtype(pipe) | |
| device = pipeline_device(pipe) | |
| cond_t = torch.from_numpy((np.array(cond_image, dtype=np.float32) / 127.5) - 1.0) | |
| cond_t = cond_t.permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=dtype) | |
| mask_t = torch.from_numpy((np.array(mask) > 127).astype(np.float32)) | |
| mask_t = mask_t.unsqueeze(0).to(device=device, dtype=dtype) | |
| masked_image_latents = build_masked_image_latents( | |
| pipe.vae, cond_t, mask_t, dtype, vae_scale_factor=pipe.vae_scale_factor | |
| ).to(device) | |
| prompt_embeds = _encode_t5(pipe, prompt, max_sequence_length).to(device) | |
| pooled_prompt_embeds = zero_pooled_projections(prompt_embeds.shape[0], device, dtype) | |
| generator = torch.Generator(device=device).manual_seed(int(seed)) | |
| return pipe( | |
| height=height, | |
| width=width, | |
| masked_image_latents=masked_image_latents, # skips the internal image*(1-mask) | |
| prompt_embeds=prompt_embeds, | |
| pooled_prompt_embeds=pooled_prompt_embeds, | |
| num_inference_steps=num_steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| max_sequence_length=max_sequence_length, | |
| ).images[0] | |
| def run_editing_inference(pipe, real_pil, text, target_polygon, other_polygons, font, | |
| num_steps, guidance_scale, seed, size=512, | |
| deterministic=True, allow_pad=True, max_sequence_length=512): | |
| """Build the condition, run inference, drop the glyph strip. | |
| Returns (scene image, polygon in scene coordinates) or (None, None). | |
| """ | |
| sample = build_sample_concat( | |
| real_pil, text, target_polygon, other_polygons, font, | |
| size=size, loose_mask=False, mask_aug=False, | |
| rng=None, deterministic=deterministic, allow_pad=allow_pad, | |
| ) | |
| if sample is None: | |
| return None, None | |
| gen = run_one_inference( | |
| pipe, sample["combined_cond_image"], sample["combined_mask"], sample["prompt"], | |
| num_steps=num_steps, guidance_scale=guidance_scale, seed=seed, | |
| max_sequence_length=max_sequence_length, | |
| ) | |
| return crop_scene_part(gen, size), sample["polygon"] | |
| def crop_with_padding(img, polygon, pad_ratio=0.0, min_pad=0): | |
| """Crop the polygon bounding box with optional padding.""" | |
| x0, y0, x1, y1 = polygon_bbox(polygon) | |
| w = max(x1 - x0, 1.0) | |
| h = max(y1 - y0, 1.0) | |
| pad_x = max(w * pad_ratio, min_pad) | |
| pad_y = max(h * pad_ratio, min_pad) | |
| img_w, img_h = img.size | |
| cx0 = max(0, int(round(x0 - pad_x))) | |
| cy0 = max(0, int(round(y0 - pad_y))) | |
| cx1 = min(img_w, int(round(x1 + pad_x))) | |
| cy1 = min(img_h, int(round(y1 + pad_y))) | |
| if cx1 <= cx0: | |
| cx1 = min(img_w, cx0 + 1) | |
| if cy1 <= cy0: | |
| cy1 = min(img_h, cy0 + 1) | |
| return img.crop((cx0, cy0, cx1, cy1)), (cx0, cy0, cx1, cy1) | |
| # ============================================================================= | |
| # Demo API: page + box + text -> crop (+ page preview) | |
| # ============================================================================= | |
| def box_to_polygon(box): | |
| x0, y0, x1, y1 = [float(v) for v in box] | |
| if x1 < x0: | |
| x0, x1 = x1, x0 | |
| if y1 < y0: | |
| y0, y1 = y1, y0 | |
| return [[x0, y0], [x1, y0], [x1, y1], [x0, y1]] | |
| def _locate_window(real_pil, target_polygon, size): | |
| """Recompute the crop window (same deterministic path as build_sample_concat).""" | |
| tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon) | |
| needed = max(int(size), int(math.ceil(max(tx1 - tx0, ty1 - ty0)))) | |
| padded, poly_p, _, pad = _pad_to_fit(real_pil, target_polygon, [], needed) | |
| crop = compute_crop_box(padded.size[0], padded.size[1], poly_p, [], | |
| target_size=size, deterministic=True) | |
| if crop is None: | |
| return None | |
| sx0, sy0, L, _ = crop | |
| return padded, poly_p, pad, sx0, sy0, L | |
| def paste_back(real_pil, gen_scene, target_polygon, size): | |
| """Paste the edited box back into the full page (preview only).""" | |
| loc = _locate_window(real_pil, target_polygon, size) | |
| if loc is None: | |
| return real_pil | |
| padded, poly_p, pad, sx0, sy0, L = loc | |
| scale = size / float(L) | |
| local = [[(px - sx0) * scale, (py - sy0) * scale] for px, py in poly_p] | |
| bx0, by0, bx1, by1 = polygon_bbox(local) | |
| gx0 = max(0, int(math.floor(bx0))) | |
| gy0 = max(0, int(math.floor(by0))) | |
| gx1 = min(gen_scene.width, int(math.ceil(bx1))) | |
| gy1 = min(gen_scene.height, int(math.ceil(by1))) | |
| if gx1 <= gx0 or gy1 <= gy0: | |
| return real_pil | |
| patch = gen_scene.crop((gx0, gy0, gx1, gy1)) | |
| ox0 = int(round(sx0 + gx0 / scale)) | |
| oy0 = int(round(sy0 + gy0 / scale)) | |
| ox1 = int(round(sx0 + gx1 / scale)) | |
| oy1 = int(round(sy0 + gy1 / scale)) | |
| Wp, Hp = padded.size | |
| ox0, oy0 = max(0, min(ox0, Wp - 1)), max(0, min(oy0, Hp - 1)) | |
| ox1, oy1 = max(ox0 + 1, min(ox1, Wp)), max(oy0 + 1, min(oy1, Hp)) | |
| patch = patch.resize((ox1 - ox0, oy1 - oy0), Image.BICUBIC) | |
| out = padded.copy() | |
| out.paste(patch, (ox0, oy0)) | |
| W0, H0 = real_pil.size | |
| return out.crop((pad, pad, pad + W0, pad + H0)) | |
| def load_font(path=DEFAULT_FONT_PATH, size=60): | |
| try: | |
| return ImageFont.truetype(path, size=size) | |
| except Exception as e: | |
| print(f"[warn] failed to load font {path}: {e}") | |
| return ImageFont.load_default() | |
| def edit_text_region(pipe, font, image, box, text, | |
| size=DEFAULT_SIZE, steps=20, guidance_scale=30.0, seed=0, | |
| crop_pad_ratio=0.0): | |
| """Page + box + new text -> dict with the cropped result and a page preview.""" | |
| image = image.convert("RGB") | |
| W, H = image.size | |
| x0, y0, x1, y1 = [float(v) for v in box] | |
| x0, x1 = max(0.0, min(x0, W - 1)), max(1.0, min(x1, W)) | |
| y0, y1 = max(0.0, min(y0, H - 1)), max(1.0, min(y1, H)) | |
| if x1 - x0 < 2 or y1 - y0 < 2: | |
| raise ValueError(f"box too small: {(x0, y0, x1, y1)}") | |
| polygon = box_to_polygon((x0, y0, x1, y1)) | |
| gen_scene, poly_local = run_editing_inference( | |
| pipe, image, text, polygon, [], font, | |
| num_steps=int(steps), guidance_scale=float(guidance_scale), seed=int(seed), | |
| size=int(size), deterministic=True, allow_pad=True, | |
| ) | |
| if gen_scene is None: | |
| raise ValueError("no valid square crop window; please draw a smaller box") | |
| crop, crop_box = crop_with_padding(gen_scene, poly_local, pad_ratio=float(crop_pad_ratio)) | |
| page = paste_back(image, gen_scene, polygon, int(size)) | |
| return { | |
| "crop": crop, | |
| "page": page, | |
| "crop_box": crop_box, | |
| "polygon_local": poly_local, | |
| "canvas": concat_canvas_size(int(size)), | |
| "prompt": generate_prompt_concat([text]), | |
| } | |