Update app.py

app.py

@@ -27,54 +27,55 @@ import torch
 import numpy as np
 from PIL import Image
 import gradio as gr
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionInpaintPipeline
 
-pipe = None
+# Lazy pipeline loader for ZeroGPU
+pipe_inpaint = None
 
-def get_pipeline():
-    global pipe
-    if pipe is None:
-        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
+def get_inpaint_pipeline():
+    global pipe_inpaint
+    if pipe_inpaint is None:
+        pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
             torch_dtype=torch.float16,
             safety_checker=None
         )
-        pipe.enable_attention_slicing()
-    return pipe
+        pipe_inpaint.enable_attention_slicing()
+    return pipe_inpaint
 
-def create_seamless_matrix(pil_img, blend_factor=0.35, mode="Horizontal Only (Landscapes/City Facades)"):
-    """
-    Crossfades edge boundaries.
-    Supports both 2D wrapping and 1D horizontal-only wrapping.
-    """
-    img = np.array(pil_img.convert("RGB")).astype(np.float32) / 255.0
-    h, w, c = img.shape
+def shift_image(pil_img, mode="Horizontal Only (Landscapes/City Facades)"):
+    """Shifts image by 50% along selected axes to position outer seams at canvas center."""
+    img = np.array(pil_img.convert("RGB"))
+    h, w, _ = img.shape
 
     if "Horizontal" in mode:
-        # Shift 50% along width (X-axis) only
         shifted = np.roll(img, w // 2, axis=1)
-
-        x = np.linspace(-1, 1, w)
-        xx = np.tile(x, (h, 1))
-
-        mask_x = np.clip(np.abs(xx) / blend_factor, 0, 1)[:, :, np.newaxis]
-        blended = shifted * mask_x + img * (1.0 - mask_x)
     else:
-        # Full 2D Shift (X and Y axes)
         shifted = np.roll(np.roll(img, h // 2, axis=0), w // 2, axis=1)
 
-        y = np.linspace(-1, 1, h)
-        x = np.linspace(-1, 1, w)
-        xx, yy = np.meshgrid(x, y)
+    return Image.fromarray(shifted)
 
-        mask_x = np.clip(np.abs(xx) / blend_factor, 0, 1)
-        mask_y = np.clip(np.abs(yy) / blend_factor, 0, 1)
-        mask = np.minimum(mask_x, mask_y)[:, :, np.newaxis]
+def create_seam_mask(width, height, seam_width_pct=0.25, mode="Horizontal Only (Landscapes/City Facades)"):
+    """
+    Creates a black-and-white mask.
+    White = Area to be inpainted by AI (center seam join zone).
+    Black = Area to keep 100% original photo intact.
+    """
+    mask = np.zeros((height, width), dtype=np.uint8)
 
-        blended = shifted * mask + img * (1.0 - mask)
+    if "Horizontal" in mode:
+        w_seam = int(width * seam_width_pct)
+        x_start = (width - w_seam) // 2
+        mask[:, x_start:x_start + w_seam] = 255
+    else:
+        w_seam = int(width * seam_width_pct)
+        h_seam = int(height * seam_width_pct)
+        x_start = (width - w_seam) // 2
+        y_start = (height - h_seam) // 2
+        mask[:, x_start:x_start + w_seam] = 255
+        mask[y_start:y_start + h_seam, :] = 255
 
-    result = (np.clip(blended, 0, 1) * 255).astype(np.uint8)
-    return Image.fromarray(result)
+    return Image.fromarray(mask)
 
 def create_tiled_preview(pil_img, mode="Horizontal Only (Landscapes/City Facades)"):
     """Generates 1x3 horizontal panorama or 3x3 grid preview."""
@@ -85,12 +86,10 @@ def create_tiled_preview(pil_img, mode="Horizontal Only (Landscapes/City Facades
     w, h = preview_img.size
 
     if "Horizontal" in mode:
-        # 1x3 Horizontal panorama test strip
         tiled = Image.new("RGB", (w * 3, h))
         for c in range(3):
             tiled.paste(preview_img, (c * w, 0))
     else:
-        # 3x3 Full 2D grid test
         tiled = Image.new("RGB", (w * 3, h * 3))
         for r in range(3):
             for c in range(3):
@@ -99,50 +98,48 @@ def create_tiled_preview(pil_img, mode="Horizontal Only (Landscapes/City Facades
     return tiled
 
 @spaces.GPU
-def generate_ai_texture(input_image, prompt, tiling_mode, steps, seed_val, strength, resolution_str, ut99_quantize):
+def generate_ai_texture(input_image, prompt, method_mode, tiling_mode, seam_width_pct, steps, seed_val, resolution_str, ut99_quantize):
     if input_image is None:
         return None, None, None, None, "⚠️ Please upload a source photo first."
 
     actual_seed = random.randint(0, 2147483647) if (seed_val == -1 or seed_val is None) else int(seed_val)
     res_val = int(resolution_str.split("x")[0])
 
-    pipeline = get_pipeline()
+    pipeline = get_inpaint_pipeline()
     pipeline.to("cuda")
     generator = torch.Generator("cuda").manual_seed(actual_seed)
 
-    # 1. Axis-aware seamless edge pre-processing
-    seamless_input = create_seamless_matrix(input_image, blend_factor=0.30, mode=tiling_mode)
-    init_img = seamless_input.resize((512, 512), Image.Resampling.LANCZOS)
+    # 1. Prepare shifted image and resolution
+    init_img = input_image.resize((512, 512), Image.Resampling.LANCZOS)
+    shifted_init = shift_image(init_img, mode=tiling_mode)
 
-    # 2. Step math guardrail
-    user_steps = max(5, int(steps))
-    user_strength = max(0.1, float(strength))
-    calculated_steps = max(int(user_steps / user_strength), user_steps + 5)
-
-    if "Horizontal" in tiling_mode:
-        axis_prompt = "horizontally seamless texture, panorama horizon, landscape game asset"
+    # 2. Build Inpainting Mask vs Full Mask
+    if "Inpainting" in method_mode:
+        seam_mask = create_seam_mask(512, 512, seam_width_pct=seam_width_pct / 100.0, mode=tiling_mode)
     else:
-        axis_prompt = "fully seamless tileable texture, 2D pattern, game asset"
+        # Full mask regenerates entire canvas
+        seam_mask = Image.new("L", (512, 512), 255)
 
-    full_prompt = prompt + ", " + axis_prompt
-    negative_prompt = "blurry, distorted, frames, borders, window glare, watermark"
+    user_steps = max(10, int(steps))
+    full_prompt = prompt + ", seamless tileable game asset, detailed continuous pattern, high quality"
+    negative_prompt = "seam line, border, frame, edge line, blur, watermark, distortion"
 
-    # 3. GPU Diffusion Execution
-    result = pipeline(
+    # 3. Perform AI Inpainting
+    inpainted_result = pipeline(
         prompt=full_prompt,
         negative_prompt=negative_prompt,
-        image=init_img,
-        strength=user_strength,
-        num_inference_steps=calculated_steps,
+        image=shifted_init,
+        mask_image=seam_mask,
+        num_inference_steps=user_steps,
         guidance_scale=7.5,
         generator=generator
     ).images[0]
 
-    # 4. Upscale to requested resolution (up to 8K)
-    final_img = result.resize((res_val, res_val), Image.Resampling.LANCZOS)
+    # 4. Upscale to target resolution (up to 8K)
+    final_img = inpainted_result.resize((res_val, res_val), Image.Resampling.LANCZOS)
     tiled_preview = create_tiled_preview(final_img, mode=tiling_mode)
 
-    # 5. Export 8-bit UT99 formats (.PCX / .BMP)
+    # 5. Export 8-bit UT99 Formats (.PCX / .BMP)
     temp_dir = tempfile.mkdtemp()
     pcx_path = os.path.join(temp_dir, f"ut99_texture_{res_val}.pcx")
     bmp_path = os.path.join(temp_dir, f"ut99_texture_{res_val}.bmp")
@@ -157,18 +154,24 @@ def generate_ai_texture(input_image, prompt, tiling_mode, steps, seed_val, stren
         final_img.save(bmp_path, format="BMP")
         display_img = final_img
 
-    status_msg = f"✅ Success | Mode: {tiling_mode} | Seed: {actual_seed} | Output: {res_val}x{res_val}"
+    status_msg = f"✅ Success | Method: {method_mode} | Seed: {actual_seed} | Output: {res_val}x{res_val}"
     return display_img, tiled_preview, pcx_path, bmp_path, status_msg
 
 # Gradio Interface
 with gr.Blocks(title="UT99 ZeroGPU AI Texture Forge") as demo:
     gr.Markdown("# ⚡ UT99 ZeroGPU AI Texture Forge")
-    gr.Markdown("Generate seamless 8-bit UT99 textures up to 8K with 1D Horizontal or 2D Full Tiling.")
+    gr.Markdown("Generate seamless 8-bit UT99 textures using Smart Seam Inpainting or Full AI Synthesis.")
 
     with gr.Row():
         with gr.Column():
             input_img = gr.Image(type="pil", label="Upload Source Photo")
-            
+
+            method_mode = gr.Radio(
+                choices=["⚡ Smart Seam Inpainting (Preserves Photo)", "🎨 Full Canvas Generation"],
+                value="⚡ Smart Seam Inpainting (Preserves Photo)",
+                label="Generation Method"
+            )
+
             tiling_mode = gr.Radio(
                 choices=["Horizontal Only (Landscapes/City Facades)", "Full 2D (Both Axes)"],
                 value="Horizontal Only (Landscapes/City Facades)",
@@ -176,19 +179,19 @@ with gr.Blocks(title="UT99 ZeroGPU AI Texture Forge") as demo:
             )
 
             prompt = gr.Textbox(
-                value="Canyon mountains horizon vista, landscape panorama texture",
-                label="AI Prompt"
+                value="Mountain ridge horizon panorama, detailed rocky landscape texture",
+                label="AI Inpainting Prompt (Describe how the seam should blend)"
+            )
+
+            seam_width_pct = gr.Slider(
+                minimum=10, maximum=50, value=25, step=5,
+                label="Inpaint Seam Width (% of image centered over join)"
             )
 
             with gr.Row():
-                steps = gr.Slider(minimum=5, maximum=30, value=15, step=1, label="Denoising Steps")
+                steps = gr.Slider(minimum=10, maximum=40, value=20, step=1, label="Denoising Steps")
                 seed_val = gr.Number(value=-1, precision=0, label="Seed (-1 = Random)")
 
-            strength = gr.Slider(
-                minimum=0.1, maximum=0.9, value=0.35, step=0.05,
-                label="AI Strength"
-            )
-
             resolution = gr.Dropdown(
                 choices=[
                     "128x128",
@@ -208,7 +211,7 @@ with gr.Blocks(title="UT99 ZeroGPU AI Texture Forge") as demo:
                 label="Convert to 8-Bit / 256 Colors (.PCX / .BMP)"
             )
 
-            btn = gr.Button("🚀 Generate AI Texture on GPU", variant="primary")
+            btn = gr.Button("🚀 Generate AI Inpainted Texture", variant="primary")
 
         with gr.Column():
             output_status = gr.Textbox(label="Status / Seed", interactive=False)
@@ -221,7 +224,10 @@ with gr.Blocks(title="UT99 ZeroGPU AI Texture Forge") as demo:
 
     btn.click(
         fn=generate_ai_texture,
-        inputs=[input_img, prompt, tiling_mode, steps, seed_val, strength, resolution, ut99_quantize],
+        inputs=[
+            input_img, prompt, method_mode, tiling_mode,
+            seam_width_pct, steps, seed_val, resolution, ut99_quantize
+        ],
         outputs=[output_texture, output_preview, pcx_file, bmp_file, output_status]
     )