LiXiY/Reference-Anomaly-Generation
0
1import argparse2import base643import os4os.system("pip uninstall -y hf-gradio")5os.system("pip uninstall -y spaces")6os.system("pip uninstall -y mcp")7os.system("pip install -r requirement.txt")8import spaces9from io import BytesIO10import gradio as gr11import torch12import torch.nn as nn13from diffusers import DDIMScheduler14from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor15from ReferenceNet import ReferenceNet16 17from inpainting_pipeline import StableDiffusionInpaintPipeline18from PIL import Image19from huggingface_hub import snapshot_download20import cv221import numpy as np22import math23 24 25class LinearResampler(nn.Module):26 def __init__(self, input_dim=1024, output_dim=1024):27 super().__init__()28 self.projector = nn.Linear(input_dim, output_dim)29 30 def forward(self, x):31 return self.projector(x)32 33 34# ===================== Attention Capture =====================35 36class CaptureAttnProcessor(nn.Module):37 def __init__(self):38 self.captured_attn_map = None39 self.cnt = 040 super().__init__()41 42 def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None):43 residual = hidden_states44 45 if attn.spatial_norm is not None:46 hidden_states = attn.spatial_norm(hidden_states, temb)47 48 input_ndim = hidden_states.ndim49 if input_ndim == 4:50 batch_size, channel, height, width = hidden_states.shape51 hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)52 53 batch_size, sequence_length, _ = (54 hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape55 )56 attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)57 58 if attn.group_norm is not None:59 hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)60 61 query = attn.to_q(hidden_states)62 63 if encoder_hidden_states is None:64 encoder_hidden_states = hidden_states65 elif attn.norm_cross:66 encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)67 68 key = attn.to_k(encoder_hidden_states)69 value = attn.to_v(encoder_hidden_states)70 71 query = attn.head_to_batch_dim(query)72 key = attn.head_to_batch_dim(key)73 value = attn.head_to_batch_dim(value)74 75 attention_probs = attn.get_attention_scores(query, key, attention_mask)76 77 if self.cnt % 3 == 0:78 if attention_probs.shape[0] > 8:79 self.captured_attn_map = attention_probs[8:, :, :].detach()80 else:81 self.captured_attn_map = attention_probs.detach()82 83 self.cnt += 184 85 hidden_states = torch.bmm(attention_probs, value)86 hidden_states = attn.batch_to_head_dim(hidden_states)87 88 hidden_states = attn.to_out[0](hidden_states)89 hidden_states = attn.to_out[1](hidden_states)90 91 if input_ndim == 4:92 hidden_states = hidden_states.transpose(-1, 2).reshape(batch_size, channel, height, width)93 94 if attn.residual_connection:95 hidden_states = hidden_states + residual96 97 hidden_states = hidden_states / attn.rescale_output_factor98 99 return hidden_states100 101 102def visualize_attention_map(hooks_dict, inpainting_mask=None, ref_image=None):103 """104 Visualize the attention map and return a PIL Image (without saving to disk).105 Uses cv2 COLORMAP_JET heatmap overlaid on the reference image.106 """107 valid_hooks = {k: v for k, v in hooks_dict.items() if v.captured_attn_map is not None}108 num_layers = len(valid_hooks)109 110 if num_layers == 0:111 print("No attention maps captured.")112 return None113 114 sorted_keys = sorted(valid_hooks.keys())115 layer_name = sorted_keys[0]116 proc = valid_hooks[layer_name]117 attention_map = proc.captured_attn_map118 119 attn_avg = attention_map.mean(dim=0).cpu().detach()120 121 split_idx = int(attn_avg.shape[0] / 2)122 height = width = int(math.sqrt(split_idx))123 124 attn_cross = attn_avg[:split_idx, split_idx:]125 126 if inpainting_mask is not None:127 mask_resized = inpainting_mask.resize((width, height), resample=Image.NEAREST)128 mask_array = np.array(mask_resized).astype(np.float32)129 if mask_array.max() > 1.0:130 mask_array = mask_array / 255.0131 mask_binary = mask_array > 0.5132 mask_flat = mask_binary.flatten()133 num_mask_pixels = int(mask_flat.sum())134 if num_mask_pixels > 0:135 attn_mask_region = attn_cross[mask_flat, :]136 map_data = attn_mask_region.mean(dim=0).reshape(height, width).numpy()137 print(f"Using {num_mask_pixels}/{len(mask_flat)} query positions from mask region")138 else:139 print("Warning: mask region is empty after binarization, falling back to full attention map")140 map_data = attn_cross.mean(dim=0).reshape(height, width).numpy()141 else:142 map_data = attn_cross.mean(dim=0).reshape(height, width).numpy()143 144 # Normalize to 0โ255145 map_data = (map_data - map_data.min()) / (map_data.max() - map_data.min() + 1e-8)146 heatmap = (map_data * 255).astype(np.uint8)147 148 # JET pseudo-color mapping149 heatmap_img = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)150 heatmap_img = cv2.cvtColor(heatmap_img, cv2.COLOR_BGR2RGB)151 152 target_size = (512, 512)153 heatmap_img = cv2.resize(heatmap_img, target_size)154 155 # Overlay onto the reference image156 if ref_image is not None:157 org_img = np.array(ref_image.resize(target_size, resample=Image.BICUBIC))158 org_img = cv2.resize(org_img, target_size)159 attn_vis = cv2.addWeighted(org_img, 0.3, heatmap_img, 0.7, 0)160 else:161 attn_vis = heatmap_img162 163 return Image.fromarray(attn_vis)164 165 166class AttentionVisualizer:167 def __init__(self, pipe):168 self.pipe = pipe169 self.hooks = {}170 171 def register_specific_layer(self):172 unet = self.pipe.unet173 target_block_idx = 3174 target_attn_idx = 2175 try:176 block = unet.up_blocks[target_block_idx]177 if hasattr(block, "attentions") and len(block.attentions) > target_attn_idx:178 attn_module = block.attentions[target_attn_idx]179 for k, transformer in enumerate(attn_module.transformer_blocks):180 target_attn = transformer.attn1181 layer_name = f"up_blocks.{target_block_idx}.attentions.{target_attn_idx}"182 hook_proc = CaptureAttnProcessor()183 target_attn.set_processor(hook_proc)184 self.hooks[layer_name] = hook_proc185 print(f"Successfully registered hook: {layer_name}")186 else:187 print(f"Error: Layer up_blocks.{target_block_idx}.attentions.{target_attn_idx} does not exist.")188 except IndexError:189 print(f"Error: up_blocks index {target_block_idx} out of range.")190 191 def reset(self):192 """Reset the capture state of all hooks. Should be called before each generation."""193 for proc in self.hooks.values():194 proc.captured_attn_map = None195 proc.cnt = 0196 197 def visualize(self, inpainting_mask=None, ref_image=None):198 """Return a PIL Image visualization of the attention map."""199 return visualize_attention_map(200 self.hooks, inpainting_mask=inpainting_mask, ref_image=ref_image201 )202 203 204# ===================== Model Wrapper =====================205 206class ReferencenetInpainting:207 def __init__(self, sd_pipe, referencenet, image_encoder_path, checkpoint_path, device):208 self.device = device209 self.image_encoder_path = image_encoder_path210 self.checkpoint_path = checkpoint_path211 self.referencenet = referencenet.to(self.device)212 self.pipe = sd_pipe.to(self.device)213 214 self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to(215 self.device, dtype=torch.float16216 )217 218 self.clip_image_processor = CLIPImageProcessor()219 self.image_proj_model = self.init_proj()220 self.load_unet_and_image_proj_and_referencenet()221 222 def init_proj(self):223 image_proj_model = LinearResampler(224 input_dim=1280,225 output_dim=self.pipe.unet.config.cross_attention_dim,226 ).to(self.device, dtype=torch.float16)227 return image_proj_model228 229 def load_unet_and_image_proj_and_referencenet(self):230 state_dict = torch.load(self.checkpoint_path, map_location="cpu")231 self.pipe.unet.load_state_dict(state_dict["unet"], strict=False)232 self.referencenet.load_state_dict(state_dict["referencenet"], strict=False)233 self.image_proj_model.load_state_dict(state_dict["image_proj"])234 235 @torch.inference_mode()236 def get_image_embeds(self, pil_image=None, clip_image_embeds=None):237 if isinstance(pil_image, Image.Image):238 pil_image = [pil_image]239 clip_image = self.clip_image_processor(images=pil_image, return_tensors="pt").pixel_values240 241 clip_image = clip_image.to(self.device, dtype=torch.float16)242 clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]243 244 image_prompt_embeds = self.image_proj_model(clip_image_embeds).to(dtype=torch.float16)245 246 uncond_clip_image_embeds = self.image_encoder(247 torch.zeros_like(clip_image), output_hidden_states=True248 ).hidden_states[-2]249 uncond_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds)250 return image_prompt_embeds, uncond_image_prompt_embeds251 252 253 def generate(254 self,255 pil_ref_image=None,256 pil_background_image=None,257 pil_mask_image=None,258 num_samples=1,259 seed=None,260 guidance_scale=7.5,261 num_inference_steps=30,262 **kwargs,263 ):264 image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(pil_image=pil_ref_image)265 bs_embed, seq_len, _ = image_prompt_embeds.shape266 image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)267 image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)268 uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)269 uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)270 generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None271 272 images = self.pipe(273 image=pil_background_image,274 mask_image=pil_mask_image,275 prompt_embeds=image_prompt_embeds,276 negative_prompt_embeds=uncond_image_prompt_embeds,277 guidance_scale=guidance_scale,278 num_inference_steps=num_inference_steps,279 generator=generator,280 referencenet=self.referencenet,281 ref_image=pil_ref_image,282 clip_image_embed=torch.cat([uncond_image_prompt_embeds, image_prompt_embeds], dim=0),283 **kwargs,284 ).images285 286 return images287 288 289# ===================== Model Setup =====================290 291parser = argparse.ArgumentParser(description="Gradio Demo")292 293allow_sd_text_encoder_patterns = ["text_encoder/config.json", "text_encoder/pytorch_model.bin"]294allow_tokenizer_patterns = ["tokenizer/*"]295allow_scheduler_patterns = ["scheduler/*"]296allow_vae_patterns = ["vae/config.json", "vae/diffusion_pytorch_model.bin"]297allow_unet_patterns = ["unet/config.json", "unet/diffusion_pytorch_model.bin"]298allow_sd_patterns = allow_sd_text_encoder_patterns + allow_tokenizer_patterns + allow_scheduler_patterns + allow_vae_patterns + allow_unet_patterns + ["model_index.json"]299 300sd_model_path = snapshot_download("stable-diffusion-v1-5/stable-diffusion-inpainting", allow_patterns=allow_sd_patterns)301ref_model_path = snapshot_download("stable-diffusion-v1-5/stable-diffusion-v1-5", allow_patterns=allow_sd_patterns)302image_encoder_path = snapshot_download("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", allow_patterns=["config.json", "pytorch_model.bin"])303checkpoint_path = snapshot_download('LiXiY/ReferenceAnomaly') + "/" + "reference_anomaly_checkponint.bin"304 305device = "cuda" if torch.cuda.is_available() else "cpu"306args = parser.parse_args()307 308noise_scheduler = DDIMScheduler(309 num_train_timesteps=1000,310 beta_start=0.00085,311 beta_end=0.012,312 beta_schedule="scaled_linear",313 clip_sample=False,314 set_alpha_to_one=False,315 steps_offset=1,316)317 318pipe = StableDiffusionInpaintPipeline.from_pretrained(319 sd_model_path,320 torch_dtype=torch.float16,321 scheduler=noise_scheduler,322 feature_extractor=None,323 safety_checker=None324)325 326referencenet = ReferenceNet.from_pretrained(ref_model_path, subfolder="unet", feature_extractor=None, safety_checker=None).to(dtype=torch.float16)327 328reference_anomaly_model = ReferencenetInpainting(pipe, referencenet, image_encoder_path, checkpoint_path, device)329 330# ===================== Register Attention Hook =====================331attention_visualizer = AttentionVisualizer(reference_anomaly_model.pipe)332attention_visualizer.register_specific_layer()333 334 335# ===================== Example Data =====================336 337CANVAS_W, CANVAS_H = 512, 512338 339EXAMPLES = [340 ("validation_images/background_image_1.png", "validation_images/ref_image_1.png", "validation_images/inpainting_mask_1.png"),341 ("validation_images/background_image_2.png", "validation_images/ref_image_2.png", "validation_images/inpainting_mask_2.png"),342 ("validation_images/background_image_3.png", "validation_images/ref_image_3.png", "validation_images/inpainting_mask_3.png"),343 ("validation_images/background_image_4.png", "validation_images/ref_image_4.png", "validation_images/inpainting_mask_4.png"),344]345 346 347# ===================== Thumbnail HTML Generation =====================348 349def img_to_b64(path, size=(120, 120)):350 img = Image.open(path).convert("RGB").resize(size, Image.LANCZOS)351 buf = BytesIO()352 img.save(buf, format="PNG")353 return base64.b64encode(buf.getvalue()).decode()354 355 356def build_examples_html():357 row_pairs = [EXAMPLES[0:2], EXAMPLES[2:4]]358 359 html = '<div class="ex-grid">'360 for row_idx, row_examples in enumerate(row_pairs):361 html += '<div class="ex-grid-row">'362 for col_idx, (bg, ref, mask) in enumerate(row_examples):363 i = row_idx * 2 + col_idx364 bg_b64 = img_to_b64(bg)365 ref_b64 = img_to_b64(ref)366 mask_b64 = img_to_b64(mask)367 html += f'''368 <div class="ex-row" onclick="(function(){{var el=document.getElementById('ex_btn_{i}');if(!el)return;var btn=el.querySelector('button')||el;btn.dispatchEvent(new MouseEvent('click',{{bubbles:true,cancelable:true}}));}})()">369 <div class="ex-label">Example {i + 1}</div>370 <div class="ex-thumbs">371 <div class="ex-thumb-wrap">372 <img src="data:image/png;base64,{bg_b64}" class="ex-thumb" draggable="false"/>373 <span class="ex-thumb-sublabel">Background</span>374 </div>375 <div class="ex-thumb-wrap">376 <img src="data:image/png;base64,{mask_b64}" class="ex-thumb" draggable="false"/>377 <span class="ex-thumb-sublabel">Mask</span>378 </div>379 <div class="ex-thumb-wrap">380 <img src="data:image/png;base64,{ref_b64}" class="ex-thumb" draggable="false"/>381 <span class="ex-thumb-sublabel">Reference</span>382 </div>383 </div>384 </div>385 '''386 html += '</div>'387 html += '</div>'388 return html389 390 391# ===================== Utility Functions =====================392 393def fit_image_to_canvas(img, canvas_w=CANVAS_W, canvas_h=CANVAS_H):394 img_rgba = img.convert("RGBA")395 img_rgba.thumbnail((canvas_w, canvas_h), Image.LANCZOS)396 canvas = Image.new("RGBA", (canvas_w, canvas_h), (0, 0, 0, 0))397 offset_x = (canvas_w - img_rgba.width) // 2398 offset_y = (canvas_h - img_rgba.height) // 2399 canvas.paste(img_rgba, (offset_x, offset_y))400 return canvas401 402 403def extract_mask_from_layers(layers, target_size):404 mask = Image.new("L", target_size, 0)405 for layer in layers:406 if layer is not None:407 layer_rgba = layer.convert("RGBA").resize(target_size)408 alpha = layer_rgba.split()[3]409 alpha_binary = alpha.point(lambda x: 255 if x > 0 else 0)410 mask = Image.composite(Image.new("L", target_size, 255), mask, alpha_binary)411 return mask412 413 414def load_example(idx):415 bg_path, ref_path, mask_path = EXAMPLES[idx]416 417 bg = Image.open(bg_path).convert("RGB").resize((CANVAS_W, CANVAS_H))418 ref_img = Image.open(ref_path).convert("RGB").resize((CANVAS_W, CANVAS_H))419 mask = Image.open(mask_path).convert("L").resize((CANVAS_W, CANVAS_H))420 421 transparent = Image.new("RGBA", (CANVAS_W, CANVAS_H), (0, 0, 0, 0))422 white_solid = Image.new("RGBA", (CANVAS_W, CANVAS_H), (255, 255, 255, 255))423 mask_layer = Image.composite(white_solid, transparent, mask)424 425 composite = Image.alpha_composite(bg.convert("RGBA"), mask_layer)426 427 editor_val = {428 "background": bg,429 "layers": [mask_layer],430 "composite": composite,431 }432 return editor_val, ref_img433 434 435def load_ex1():436 return load_example(0)437 438def load_ex2():439 return load_example(1)440 441def load_ex3():442 return load_example(2)443 444def load_ex4():445 return load_example(3)446 447 448# ===================== Generation Function (also returns attention map) =====================449@spaces.GPU450def run_local(base, ref):451 if base is None or ref is None:452 return None, None, gr.update(visible=False)453 454 target_size = (CANVAS_W, CANVAS_H)455 pil_ref = ref.convert("RGB").resize(target_size)456 457 if not isinstance(base, dict):458 return None, None, gr.update(visible=False)459 460 bg_pil = base.get("background")461 layers = base.get("layers", [])462 463 if bg_pil is None:464 return None, None, gr.update(visible=False)465 466 pil_bg = bg_pil.convert("RGB").resize(target_size)467 pil_mask = extract_mask_from_layers(layers, target_size)468 469 if pil_mask.getextrema() == (0, 0):470 error_html = """471 <div class="error-overlay" style="472 position: fixed; top: 0; left: 0; width: 100%; height: 100%;473 background: rgba(0,0,0,0.5); display: flex; justify-content: center;474 align-items: center; z-index: 9999;475 ">476 <div style="477 background: white; padding: 30px; border-radius: 10px;478 text-align: center; font-size: 18px; box-shadow: 0 0 15px rgba(0,0,0,0.3);479 ">480 <p style="color: red; margin-bottom: 20px;">481 โ ๏ธ Please draw the anomaly region (mask) on the background image first, or click an example!482 </p>483 <button onclick="this.closest('.error-overlay').remove()"484 style="padding: 8px 20px; cursor: pointer; border: none;485 background: #eee; border-radius: 5px;">486 OK487 </button>488 </div>489 </div>490 """491 return None, None, gr.update(value=error_html, visible=True)492 493 # Reset attention capture before generation494 attention_visualizer.reset()495 496 generated_images = reference_anomaly_model.generate(497 pil_ref_image=pil_ref,498 pil_background_image=pil_bg,499 pil_mask_image=pil_mask,500 num_samples=1,501 guidance_scale=7.5,502 num_inference_steps=25,503 seed=42,504 )505 506 result_img = generated_images[0].resize(target_size)507 508 # Generate attention map visualization509 attn_img = attention_visualizer.visualize(510 inpainting_mask=pil_mask,511 ref_image=pil_ref,512 )513 514 return result_img, attn_img, gr.update(visible=False)515 516 517# ===================== Combined Client JS (resize + force English) =====================518# KEY FIX: merge both JS functions into a SINGLE function body instead of519# concatenating two separate function expressions.520 521COMBINED_JS = """522function() {523 /* ===== Client-Side Instant Resize ===== */524 var MAX_W = """ + str(CANVAS_W) + """;525 var MAX_H = """ + str(CANVAS_H) + """;526 527 function resizeInBrowser(file) {528 return new Promise(function(resolve) {529 var reader = new FileReader();530 reader.onload = function(e) {531 var img = new Image();532 img.onload = function() {533 if (img.width <= MAX_W && img.height <= MAX_H) {534 resolve(null);535 return;536 }537 var ratio = Math.min(MAX_W / img.width, MAX_H / img.height);538 var c = document.createElement('canvas');539 c.width = Math.round(img.width * ratio);540 c.height = Math.round(img.height * ratio);541 c.getContext('2d').drawImage(img, 0, 0, c.width, c.height);542 c.toBlob(function(blob) {543 resolve(blob ? new File([blob], file.name, {type: 'image/png'}) : null);544 }, 'image/png');545 };546 img.onerror = function() { resolve(null); };547 img.src = e.target.result;548 };549 reader.onerror = function() { resolve(null); };550 reader.readAsDataURL(file);551 });552 }553 554 function hookInput(inp) {555 if (inp._resizeHooked) return;556 inp._resizeHooked = true;557 inp._skipResize = false;558 559 inp.addEventListener('change', function(e) {560 if (inp._skipResize) { inp._skipResize = false; return; }561 562 var file = inp.files && inp.files[0];563 if (!file || !file.type || file.type.indexOf('image/') !== 0) return;564 565 e.stopImmediatePropagation();566 e.stopPropagation();567 568 resizeInBrowser(file).then(function(resized) {569 if (resized) {570 var dt = new DataTransfer();571 dt.items.add(resized);572 inp.files = dt.files;573 }574 inp._skipResize = true;575 inp.dispatchEvent(new Event('change', {bubbles: true}));576 });577 }, true);578 }579 580 function scan() {581 var inputs = document.querySelectorAll('.input-row input[type="file"]');582 for (var i = 0; i < inputs.length; i++) hookInput(inputs[i]);583 }584 585 new MutationObserver(scan).observe(document.body, {childList: true, subtree: true});586 scan();587 588 /* ===== Force English UI Labels ===== */589 var zh2en = {590 'ๅฐๅพๅๆๆพๅฐๆญคๅคๆ็นๅปไธไผ ': 'Drag image here or click to upload',591 'ๆๆพๆไปถๅฐ่ฟ้': 'Drag file here',592 '็นๅปไธไผ ': 'Click to upload',593 'ๆ็นๅปไธไผ ': 'or click to upload',594 'ไธไผ ๅพ็': 'Upload image',595 '็ฒ่ดดๅพ็ๆURL': 'Paste image or URL',596 'ๆธ
็ฉบ': 'Clear',597 '็ผ่พ': 'Edit',598 'ๆค้': 'Undo',599 '้ๅ': 'Redo',600 '็ผฉๆพ': 'Zoom',601 '็ป็ฌ': 'Brush',602 'ๆฉก็ฎๆฆ': 'Eraser',603 'ๆธ
้คๅพๅฑ': 'Clear layers',604 'ๅพๅ็ผ่พๅจ': 'Image Editor',605 '็ๆ': 'Generate',606 'ๆญฃๅจ่ฟ่ก...': 'Running...',607 'ๆไบค': 'Submit',608 };609 610 function translateNode(node) {611 if (node.nodeType === Node.TEXT_NODE) {612 var text = node.textContent;613 for (var zh in zh2en) {614 if (text.indexOf(zh) !== -1) {615 text = text.split(zh).join(zh2en[zh]);616 }617 }618 if (text !== node.textContent) node.textContent = text;619 }620 }621 622 function walkAndTranslate(root) {623 var walker = document.createTreeWalker(root, NodeFilter.SHOW_TEXT, null, false);624 var node;625 while (node = walker.nextNode()) translateNode(node);626 }627 628 function translateAttributes(root) {629 root.querySelectorAll('[placeholder]').forEach(function(el) {630 var ph = el.getAttribute('placeholder');631 for (var zh in zh2en) {632 if (ph.indexOf(zh) !== -1) ph = ph.split(zh).join(zh2en[zh]);633 }634 el.setAttribute('placeholder', ph);635 });636 root.querySelectorAll('[title]').forEach(function(el) {637 var t = el.getAttribute('title');638 for (var zh in zh2en) {639 if (t.indexOf(zh) !== -1) t = t.split(zh).join(zh2en[zh]);640 }641 el.setAttribute('title', t);642 });643 }644 645 function runTranslate() {646 walkAndTranslate(document.body);647 translateAttributes(document.body);648 }649 650 var translateObserver = new MutationObserver(function(mutations) {651 for (var m = 0; m < mutations.length; m++) {652 var added = mutations[m].addedNodes;653 for (var n = 0; n < added.length; n++) {654 if (added[n].nodeType === Node.ELEMENT_NODE) {655 walkAndTranslate(added[n]);656 translateAttributes(added[n]);657 }658 }659 }660 });661 662 translateObserver.observe(document.body, { childList: true, subtree: true });663 runTranslate();664 setInterval(runTranslate, 2000);665}666"""667 668 669# ===================== Gradio UI =====================670 671with gr.Blocks(css="""672 .input-row {673 overflow: visible !important;674 }675 676 .input-row .gr-image-editor {677 overflow: hidden !important;678 }679 .input-row .gr-image-editor .image-container,680 .input-row .gr-image-editor .canvas-container,681 .input-row .gr-image-editor canvas {682 max-width: 100% !important;683 max-height: 100% !important;684 object-fit: contain !important;685 }686 687 .ex-section-header {688 display: flex;689 align-items: center;690 gap: 10px;691 margin: 28px 0 14px 0;692 justify-content: center;693 }694 .ex-section-header::before {695 content: '';696 flex: 1;697 height: 1px;698 max-width: 180px;699 background: #e5e7eb;700 }701 .ex-section-header::after {702 content: '';703 flex: 1;704 height: 1px;705 max-width: 180px;706 background: #e5e7eb;707 }708 709 .ex-container {710 display: flex;711 flex-direction: column;712 align-items: center;713 padding-bottom: 20px;714 }715 716 .ex-grid {717 display: flex;718 flex-direction: column;719 align-items: center;720 gap: 12px;721 padding-bottom: 20px;722 }723 .ex-grid-row {724 display: flex;725 gap: 20px;726 justify-content: center;727 flex-wrap: wrap;728 }729 730 .ex-row {731 display: flex;732 align-items: center;733 gap: 20px;734 padding: 14px 28px;735 border: 2px solid #e5e7eb;736 border-radius: 12px;737 cursor: pointer;738 transition: all 0.25s ease;739 background: #ffffff;740 user-select: none;741 width: fit-content;742 }743 .ex-row:hover {744 border-color: #3b82f6;745 background: #f0f7ff;746 box-shadow: 0 4px 18px rgba(59, 130, 246, 0.15);747 transform: translateY(-2px);748 }749 .ex-row:active {750 transform: translateY(0);751 box-shadow: 0 2px 8px rgba(59, 130, 246, 0.2);752 }753 754 .ex-label {755 font-weight: 700;756 font-size: 15px;757 min-width: 62px;758 color: #1e40af;759 letter-spacing: 0.02em;760 }761 762 .ex-thumbs {763 display: flex;764 gap: 14px;765 }766 767 .ex-thumb-wrap {768 display: flex;769 flex-direction: column;770 align-items: center;771 gap: 6px;772 }773 774 .ex-thumb {775 width: 110px;776 height: 110px;777 object-fit: cover;778 border-radius: 8px;779 border: 2px solid #e5e7eb;780 transition: all 0.25s ease;781 pointer-events: none;782 }783 .ex-row:hover .ex-thumb {784 border-color: #93c5fd;785 }786 787 .ex-thumb-sublabel {788 font-size: 12px;789 color: #6b7280;790 font-weight: 500;791 }792""", js=COMBINED_JS) as demo:793 794 gr.Markdown(795 "<h1 style='text-align: center;'>Reference-Based Anomaly Image Generation</h1>"796 "<h3 style='text-align: center;'>Generate anomaly images similar to the reference anomaly on normal images</h3>"797 "<h3 style='text-align: center;'>Github: https://github.com/huan-yin/reference_anomaly_generation</h3>"798 )799 gr.Markdown(800 """801 **Instructions:**802 1. Upload a background image (normal object), then use the brush tool below the image to mark the region where you want to generate an anomaly (mask), and upload a reference image (reference anomaly).803 2. Or click any row of thumbnails in the "Examples" section below to automatically load a background + mask + reference image.804 3. Click the "Generate" button, and the result will be displayed below.805 """806 )807 808 with gr.Row(elem_classes="input-row"):809 base = gr.ImageEditor(810 label="Background Image (Normal Object)",811 type="pil",812 width=420,813 height=450,814 canvas_size=(CANVAS_W, CANVAS_H),815 sources=["upload"],816 brush=gr.Brush(817 default_size=15,818 default_color="#FFFFFF",819 color_mode="fixed",820 colors=["#FFFFFF"],821 ),822 )823 ref = gr.Image(824 label="Reference Image (Reference Anomaly)",825 sources=["upload"],826 type="pil",827 width=420,828 height=380,829 )830 831 with gr.Row():832 gen_btn = gr.Button("Generate", variant="primary")833 834 # ==================== Generation Result + Attention Map Side-by-Side ====================835 with gr.Row():836 with gr.Column(scale=1):837 output_image = gr.Image(838 label="Generated Result",839 interactive=False,840 )841 with gr.Column(scale=1):842 attention_map_output = gr.Image(843 label="Attention Map (Attention Visualization for Reference Anomaly)",844 interactive=False,845 )846 847 with gr.Row():848 error_dialog = gr.HTML(visible=False)849 850 gr.HTML('<div class="ex-section-header"><span style="font-weight:700;font-size:16px;color:#374151;">Examples (click to load background + mask + reference image)</span></div>')851 852 ex_btn0 = gr.Button("Example 1", visible=False, elem_id="ex_btn_0")853 ex_btn1 = gr.Button("Example 2", visible=False, elem_id="ex_btn_1")854 ex_btn2 = gr.Button("Example 3", visible=False, elem_id="ex_btn_2")855 ex_btn3 = gr.Button("Example 4", visible=False, elem_id="ex_btn_3")856 857 gr.HTML('<div class="ex-container">' + build_examples_html() + '</div>')858 859 # ==================== Event Bindings ====================860 861 ex_btn0.click(fn=load_ex1, outputs=[base, ref])862 ex_btn1.click(fn=load_ex2, outputs=[base, ref])863 ex_btn2.click(fn=load_ex3, outputs=[base, ref])864 ex_btn3.click(fn=load_ex4, outputs=[base, ref])865 866 gen_btn.click(867 fn=run_local,868 inputs=[base, ref],869 outputs=[output_image, attention_map_output, error_dialog],870 )871 872demo.launch(server_name="0.0.0.0", server_port=7860)873