Upload app.py

app.py

@@ -1,61 +1,62 @@
 import os
 import sys
 import requests
-import io  # 内存缓冲
+import io  # Memory buffer
 
-# Spaces环境配置
+# Spaces environment configuration
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 
 import time
 import random
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 from diffusers import StableDiffusionInpaintPipelineLegacy, DDIMScheduler, AutoencoderKL, DPMSolverMultistepScheduler
-from huggingface_hub import hf_hub_url, login  # hf_hub_url 用于生成云 URL
+from huggingface_hub import hf_hub_url, login  # hf_hub_url for generating cloud URL
 import gradio as gr
 
-# 尝试导入 Anomagic（如果有 ip_adapter 模块）
+# Attempt to import Anomagic (if ip_adapter module exists)
 try:
     from ip_adapter.ip_adapter_anomagic import Anomagic
+
     HAS_ANOMAGIC = True
 except ImportError:
     HAS_ANOMAGIC = False
-    print("Anomagic 未导入，将使用基础 Inpainting")
+    print("Anomagic not imported, will use basic Inpainting")
 
-# 获取当前脚本的绝对路径（解决路径问题）
+# Get the absolute path of the current script (to solve path issues)
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
 class SingleAnomalyGenerator:
     def __init__(self, device="cuda:0"):
-        # 自动检测GPU并设置dtype
+        # Auto-detect GPU and set dtype
         if torch.cuda.is_available() and "cuda" in device:
             self.device = torch.device(device)
             self.dtype = torch.float16
-            print(f"使用GPU: {device}, dtype: {self.dtype}")
+            print(f"Using GPU: {device}, dtype: {self.dtype}")
         else:
             self.device = torch.device("cpu")
             self.dtype = torch.float32
-            print(f"使用CPU, dtype: {self.dtype}")
+            print(f"Using CPU, dtype: {self.dtype}")
 
         self.anomagic_model = None
-        self.pipe = None  # 保存 pipe 以复用
+        self.pipe = None  # Save pipe for reuse
         self.clip_vision_model = None
         self.clip_image_processor = None
-        self.ip_ckpt_path = None  # 内存中 IP 权重 state_dict
-        self.att_ckpt_path = None  # 内存中 ATT 权重 state_dict
+        self.ip_ckpt_path = None  # IP weights state_dict in memory
+        self.att_ckpt_path = None  # ATT weights state_dict in memory
 
     def load_models(self):
         """Load models with official CLIP"""
-        print("正在加载VAE...")
+        print("Loading VAE...")
         from diffusers import AutoencoderKL
         vae = AutoencoderKL.from_pretrained(
             "stabilityai/sd-vae-ft-mse",
             torch_dtype=self.dtype
         ).to(self.device)
 
-        print("正在加载基础模型...")
+        print("Loading base model...")
         from diffusers import StableDiffusionInpaintPipelineLegacy, DDIMScheduler, DPMSolverMultistepScheduler
 
         noise_scheduler = DDIMScheduler(
@@ -80,7 +81,7 @@ class SingleAnomalyGenerator:
 
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
 
-        print("正在加载CLIP图像编码器...")
+        print("Loading CLIP image encoder...")
         from transformers import CLIPVisionModel, CLIPImageProcessor
         self.clip_vision_model = CLIPVisionModel.from_pretrained(
             "openai/clip-vit-large-patch14",
@@ -88,38 +89,39 @@ class SingleAnomalyGenerator:
         ).to(self.device)
         self.clip_image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14")
 
-        print("所有模型加载完成！")
+        print("All models loaded!")
 
-        # 加载权重（从云仓库下载到内存，避免任何磁盘使用）
-        print("正在加载权重到内存...")
+        # Load weights (download from cloud repo to memory, avoid any disk usage)
+        print("Loading weights into memory...")
         weight_files = [
             ("checkpoint/ip_adapter_0.bin", "ip_ckpt_path"),
             ("checkpoint/att.bin", "att_ckpt_path")
         ]
         for filename, attr_name in weight_files:
             try:
-                # 生成云 URL（public 仓库，无需 token）
+                # Generate cloud URL (public repo, no token needed)
                 repo_id = "yuxinjiang11/Anomagic_model"
                 url = hf_hub_url(repo_id=repo_id, filename=filename, repo_type="model")
 
-                # 动态设置属性（或用if判断显式赋值）
+                # Dynamically set attribute (or use if to assign explicitly)
                 if attr_name == "ip_ckpt_path":
                     self.ip_ckpt_path = url
                 elif attr_name == "att_ckpt_path":
                     self.att_ckpt_path = url
 
-                print(f"权重文件路径: {filename} -> {url}")
+                print(f"Weight file path: {filename} -> {url}")
             except Exception as e:
-                raise FileNotFoundError(f"无法获取权重文件路径 {filename}: {str(e)}")
+                raise FileNotFoundError(f"Unable to get weight file path {filename}: {str(e)}")
 
-        # 如果有 Anomagic，加载权重到模型
+        # If Anomagic is available, load weights into the model
         if HAS_ANOMAGIC:
-            print("初始化 Anomagic 模型...")
-            self.anomagic_model = Anomagic(self.pipe, self.clip_vision_model, self.ip_ckpt_path, self.att_ckpt_path, self.device)
+            print("Initializing Anomagic model...")
+            self.anomagic_model = Anomagic(self.pipe, self.clip_vision_model, self.ip_ckpt_path, self.att_ckpt_path,
+                                           self.device)
         else:
-            print("无 Anomagic，使用基础 Pipe。")
+            print("No Anomagic, using basic Pipe.")
 
-        print("模型加载完成！")
+        print("Model loading complete!")
 
     def generate_single_image(self, normal_image, reference_image, mask, mask_0, prompt, num_inference_steps=50,
                               ip_scale=0.3, seed=42, strength=0.3):
@@ -148,10 +150,10 @@ class SingleAnomalyGenerator:
         print(f"Generating with seed {seed}...")
         torch.manual_seed(seed)
 
-        # 如果有 Anomagic，用它生成；否则基础 Inpainting
+        # If Anomagic is available, use it to generate; otherwise basic Inpainting
         if HAS_ANOMAGIC and self.anomagic_model:
             # generator = torch.Generator(device=self.device).manual_seed(seed)
-            # 假设 Anomagic.generate 支持参数（调整根据实际）
+            # Assume Anomagic.generate supports parameters (adjust based on actual)
             generated_image = self.anomagic_model.generate(
                 pil_image=reference_image,
                 num_samples=1,
@@ -165,10 +167,10 @@ class SingleAnomalyGenerator:
                 # generator=generator
             )[0]
         else:
-            # 基础 Inpainting
+            # Basic Inpainting
             # generator = torch.Generator(device=self.device).manual_seed(seed)
             if mask is None:
-                mask = Image.new('L', target_size, 255)  # 全白 mask
+                mask = Image.new('L', target_size, 255)  # Full white mask
             generated_image = self.pipe(
                 prompt=prompt,
                 image=normal_image,
@@ -181,50 +183,66 @@ class SingleAnomalyGenerator:
         return generated_image
 
 
-# 全局 generator 和加载状态
+# Global generator and load status
 generator = None
 load_status = {"loaded": False, "error": None}
 
 
 def load_generator():
-    """Gradio 加载函数：首次运行时加载模型"""
+    """Background load function: Automatically load model on startup"""
     global generator, load_status
 
     if load_status["loaded"]:
-        return "模型已加载完成！"
+        return "Models loaded!"
 
     if load_status["error"]:
-        return f"之前加载失败: {load_status['error']}"
+        return f"Previous load failed: {load_status['error']}"
 
     try:
+        print("Starting background model load...")
         generator = SingleAnomalyGenerator()
         generator.load_models()
         load_status["loaded"] = True
-        return "模型加载完成！现在可以生成图像。"
+        print("Background model load complete!")
+        return "Model loading complete! You can now generate images."
     except Exception as e:
         load_status["error"] = str(e)
-        error_msg = f"模型加载失败: {str(e)}"
+        error_msg = f"Model loading failed: {str(e)}"
         print(error_msg)
         import traceback
         print(traceback.format_exc())
         return error_msg
 
 
+def generate_random_mask(size=(512, 512), num_blobs=3, blob_size_range=(50, 150)):
+    """Generate random mask: Create several random blobs as anomaly areas"""
+    mask = Image.new('L', size, 0)  # Black background
+    draw = ImageDraw.Draw(mask)
+    for _ in range(num_blobs):
+        x = random.randint(0, size[0])
+        y = random.randint(0, size[1])
+        width = random.randint(*blob_size_range)
+        height = random.randint(*blob_size_range)
+        # Draw elliptical blobs
+        draw.ellipse([x - width // 2, y - height // 2, x + width // 2, y + height // 2], fill=255)
+    return mask
+
+
 def generate_anomaly(normal_img, reference_img, mask_img, mask_0_img, prompt, strength, ip_scale, steps, seed):
-    """核心生成函数：Gradio 调用（支持两个mask）"""
+    """Core generation function: Called by Gradio (supports two masks)"""
     global generator
 
     if not load_status["loaded"]:
-        return None, "请先点击 '加载模型' 按钮初始化。"
+        return None, "Please wait for model loading to complete."
 
     if normal_img is None or reference_img is None or not prompt.strip():
-        return None, "请上传正常图片、参考图片，并输入提示文本。"
+        return None, "Please upload normal image, reference image, and enter prompt text."
 
     if mask_img is None:
-        return None, "请上传normal image的mask图片。"
+        return None, "Please upload or generate mask image for normal image."
 
     try:
-        # 设置种子
+        # Set seed
         random.seed(seed)
         np.random.seed(seed)
         torch.manual_seed(seed)
@@ -241,61 +259,165 @@ def generate_anomaly(normal_img, reference_img, mask_img, mask_0_img, prompt, st
             strength=strength
         )
 
-        return generated_img, f"生成成功！种子: {seed}, 步数: {steps}"
+        return generated_img, f"Generation successful! Seed: {seed}, Steps: {steps}"
 
     except Exception as e:
-        error_msg = f"生成错误: {str(e)}"
+        error_msg = f"Generation error: {str(e)}"
         print(error_msg)
         import traceback
         print(traceback.format_exc())
         return None, error_msg
 
 
+# Predefined anomaly examples (using local image paths; assume images are in examples/ folder in the same directory as the script)
+EXAMPLE_PAIRS = [
+    {
+        "normal": "examples/normal_leather.png",  # Your local normal gear image
+        "reference": "examples/reference_leather.png",  # Your local rusty gear reference image
+        "mask": "examples/normal_mask_leather.png",  # Your local mask for normal gear
+        "mask_0": "examples/ref_mask_leather.png",  # Your local mask for reference gear
+        "prompt": "Bagel has a crack running across its surface.",
+        "strength": 0.6,
+        "ip_scale": 0.1,
+        "steps": 20,
+        "seed": 42,
+        "description": "Bagel has a crack running across its surface."
+    },
+    {
+        "normal": "examples/normal_candle.JPG",  # Your local normal gear image
+        "reference": "examples/reference_candle.png",  # Your local rusty gear reference image
+        "mask": "examples/normal_mask_candle.png",  # Your local mask for normal gear
+        "mask_0": "examples/ref_mask_candle.png",  # Your local mask for reference gear
+        "prompt": "Chocolate - chip cookie has a chunk - missing defect with exposed inner texture. ",
+        "strength": 0.6,
+        "ip_scale": 0.1,
+        "steps": 20,
+        "seed": 42,
+        "description": "Chocolate - chip cookie has a chunk - missing defect with exposed inner texture. "
+    },
+    {
+        "normal": "examples/normal_apple.png",  # Your local normal gear image
+        "reference": "examples/reference_apple.png",  # Your local rusty gear reference image
+        "mask": "examples/normal_mask_apple.jpg",  # Your local mask for normal gear
+        "mask_0": "examples/ref_mask_apple.png",  # Your local mask for reference gear
+        "prompt": "Wood surface has holes with rough - edged circular openings.",
+        "strength": 0.6,
+        "ip_scale": 0.1,
+        "steps": 20,
+        "seed": 42,
+        "description": "Wood surface has holes with rough - edged circular openings."
+    }
+]
+
+
+def load_example(idx):
+    """Load example: Load images from local path, generate random mask if not provided, and set UI"""
+    if idx >= len(EXAMPLE_PAIRS):
+        return None, None, None, None, EXAMPLE_PAIRS[idx]["prompt"], EXAMPLE_PAIRS[idx]["strength"], EXAMPLE_PAIRS[idx][
+            "ip_scale"], EXAMPLE_PAIRS[idx]["steps"], EXAMPLE_PAIRS[idx][
+            "seed"], f"Example {idx + 1}: {EXAMPLE_PAIRS[idx]['description']}"
+
+    ex = EXAMPLE_PAIRS[idx]
+    try:
+        # Load normal image
+        normal_img = Image.open(ex["normal"]).convert('RGB')
+
+        # Load reference image
+        reference_img = Image.open(ex["reference"]).convert('RGB')
+
+        # Load or generate normal mask
+        if ex["mask"] is not None:
+            mask_img = Image.open(ex["mask"]).convert('L')
+        else:
+            mask_img = generate_random_mask()
+
+        # Load or generate reference mask (mask_0)
+        if ex["mask_0"] is not None:
+            mask_0_img = Image.open(ex["mask_0"]).convert('L')
+        else:
+            mask_0_img = generate_random_mask()
+
+        return normal_img, reference_img, mask_img, mask_0_img, ex["prompt"], ex["strength"], ex["ip_scale"], ex[
+            "steps"], ex["seed"], f"Example {idx + 1}: {ex['description']} loaded!"
+    except Exception as e:
+        error_msg = f"Example loading failed: {str(e)} (Check if local image paths are correct)"
+        print(error_msg)
+        # Fallback to placeholder images and random masks
+        normal_img = Image.new('RGB', (512, 512), color='gray')
+        reference_img = Image.new('RGB', (512, 512), color='blue')
+        mask_img = generate_random_mask()
+        mask_0_img = generate_random_mask()
+        return normal_img, reference_img, mask_img, mask_0_img, ex["prompt"], ex["strength"], ex["ip_scale"], ex[
+            "steps"], ex["seed"], error_msg
+
+
+# Automatically load model on startup
+load_generator()
+
 # Gradio UI
-with gr.Blocks(title="Anomagic 异常图像生成器") as demo:
-    gr.Markdown("# Anomagic: 单异常图像生成 Demo")
+with gr.Blocks(title="Anomagic Anomaly Image Generator",
+               theme=gr.themes.Soft()) as demo:  # Use Soft theme for beautification
+    gr.Markdown("# Anomagic: Single Anomaly Image Generation Demo")
     gr.Markdown(
-        "上传正常图片、参考图片、normal mask和reference mask（白色区域为修复/异常生成区），输入提示，调整参数，一键生成合成异常图像。首次使用需加载模型（几分钟）。")
+        "Upload normal image, reference image, normal mask and reference mask (white areas are for inpainting/anomaly generation), enter prompt, adjust parameters, and generate synthetic anomaly images with one click. Model is loaded in the background.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            normal_img = gr.Image(type="pil", label="正常图片 (Normal Image)")
-            reference_img = gr.Image(type="pil", label="参考图片 (Reference Image)")
+            normal_img = gr.Image(type="pil", label="Normal Image", height=300)  # Limit height
+            reference_img = gr.Image(type="pil", label="Reference Image", height=300)
+
+            with gr.Row():  # Mask row: Add buttons
+                mask_img = gr.Image(type="pil", label="Normal Image Mask (white for anomaly generation area)",
+                                    height=300, tool="sketch")  # Add sketch tool
+                gr.Button("Generate Random Normal Mask").click(lambda: generate_random_mask(), outputs=mask_img)
 
-            mask_img = gr.Image(type="pil", label="Normal Image Mask (白色为异常生成区域)")
-            mask_0_img = gr.Image(type="pil", label="Reference Image Mask (mask_0)")
+            mask_0_img = gr.Image(type="pil", label="Reference Image Mask (mask_0)", height=300,
+                                  tool="sketch")  # Add sketch tool
+            gr.Button("Generate Random Reference Mask").click(lambda: generate_random_mask(), outputs=mask_0_img)
 
-            prompt = gr.Textbox(label="提示文本 (Prompt)",
+            prompt = gr.Textbox(label="Prompt Text",
                                 placeholder="e.g., a broken machine part with rust and cracks")
 
         with gr.Column(scale=1):
-            strength = gr.Slider(0.1, 1.0, value=0.5, label="去噪强度 (Strength)")
-            ip_scale = gr.Slider(0, 2.0, value=0.3, step=0.1, label="IP 适配器缩放 (IP Scale)")
-            steps = gr.Slider(10, 100, value=20, step=5, label="推理步数 (Steps)")
-            seed = gr.Slider(0, 2 ** 32 - 1, value=42, step=1, label="随机种子 (Seed)")
+            strength = gr.Slider(0.1, 1.0, value=0.5, label="Denoising Strength")
+            ip_scale = gr.Slider(0, 2.0, value=0.3, step=0.1, label="IP Adapter Scale")
+            steps = gr.Slider(10, 100, value=20, step=5, label="Inference Steps")
+            seed = gr.Slider(0, 2 ** 32 - 1, value=42, step=1, label="Random Seed")
 
     with gr.Row():
-        load_btn = gr.Button("加载模型 (Load Models)", variant="secondary")
-        generate_btn = gr.Button("生成图像 (Generate)", variant="primary")
+        generate_btn = gr.Button("Generate Image", variant="primary", size="lg")  # Enlarge button
 
-    output_img = gr.Image(type="pil", label="生成的异常图像 (Generated Anomaly)")
-    status = gr.Textbox(label="状态 (Status)", interactive=False)
+    output_img = gr.Image(type="pil", label="Generated Anomaly Image", height=400)
+    status = gr.Textbox(label="Status", interactive=False)
 
-    # 事件绑定（修复重复输出问题）
-    load_btn.click(load_generator, outputs=status)
+    # Event bindings
     generate_btn.click(
         generate_anomaly,
         inputs=[normal_img, reference_img, mask_img, mask_0_img, prompt, strength, ip_scale, steps, seed],
-        outputs=[output_img, status]  # 修复重复绑定问题
+        outputs=[output_img, status]
     )
 
-    # 清理缓存按钮（简化，因为现在无持久下载）
+    # Examples section
+    gr.Markdown("## Examples")
+    gr.Markdown(
+        "Click the buttons below to load predefined examples for quick testing. After loading, click 'Generate Image' to view the anomaly synthesis result.")
+    with gr.Row():
+        for i in range(len(EXAMPLE_PAIRS)):
+            with gr.Column():
+                ex_btn = gr.Button(f"Example {i + 1}: {EXAMPLE_PAIRS[i]['description']}", variant="secondary")
+                ex_btn.click(load_example, inputs=gr.State(i),
+                             outputs=[normal_img, reference_img, mask_img, mask_0_img, prompt, strength, ip_scale,
+                                      steps, seed, status])
+
+
+    # Clear cache button
     def clear_cache():
         global load_status
         load_status = {"loaded": False, "error": None}
-        return "缓存已清理，请重新加载模型。"
+        return "Cache cleared, please restart the app to reload the model."
+
 
-    clear_btn = gr.Button("清理缓存 (Clear Cache)", variant="stop")
+    clear_btn = gr.Button("Clear Cache", variant="stop")
     clear_btn.click(clear_cache, outputs=status)
 
 if __name__ == "__main__":