Update app.py

app.py

@@ -1,8 +1,9 @@
 import os
 import uuid
-import html
+import math
 import gradio as gr
 import spaces
+import numpy as np
 from PIL import Image
 from image_gen_aux import UpscaleWithModel
 
@@ -11,10 +12,8 @@ from image_gen_aux import UpscaleWithModel
 # ---------------------------------
 
 BASE_TMP_DIR = "/tmp/image_enhancer"
-ORIGINAL_DIR = os.path.join(BASE_TMP_DIR, "original")
 ENHANCED_DIR = os.path.join(BASE_TMP_DIR, "enhanced")
 
-os.makedirs(ORIGINAL_DIR, exist_ok=True)
 os.makedirs(ENHANCED_DIR, exist_ok=True)
 
 # ---------------------------------
@@ -38,145 +37,6 @@ RATIO_MAP = {
     "3:2": (3, 2),
 }
 
-# ---------------------------------
-# Frontend JS/CSS
-# ---------------------------------
-
-HEAD = """
-<style>
-.preview-shell {
-    background: #1f1f1f;
-    border: 1px solid #3a3a3a;
-    border-radius: 10px;
-    padding: 14px;
-}
-.preview-toolbar {
-    display: flex;
-    flex-wrap: wrap;
-    gap: 10px;
-    align-items: center;
-    margin-bottom: 12px;
-}
-.preview-btn {
-    background: #2a2a2a;
-    color: white;
-    border: 1px solid #4a4a4a;
-    border-radius: 8px;
-    padding: 8px 12px;
-    cursor: pointer;
-    font-size: 14px;
-}
-.preview-btn:hover {
-    background: #343434;
-}
-.preview-check {
-    display: inline-flex;
-    align-items: center;
-    gap: 8px;
-    color: white;
-    font-size: 14px;
-}
-.preview-slider-wrap {
-    display: inline-flex;
-    align-items: center;
-    gap: 8px;
-    color: white;
-    font-size: 14px;
-}
-.preview-stage {
-    position: relative;
-    width: 100%;
-    min-height: 320px;
-    background: #181818;
-    border: 1px solid #333;
-    border-radius: 10px;
-    overflow: hidden;
-    display: flex;
-    align-items: center;
-    justify-content: center;
-}
-.preview-base {
-    display: block;
-    width: 100%;
-    height: auto;
-    max-height: 70vh;
-    object-fit: contain;
-}
-.preview-overlay {
-    position: absolute;
-    inset: 0;
-    width: 100%;
-    height: 100%;
-    object-fit: contain;
-    pointer-events: none;
-}
-</style>
-
-<script>
-function setPreviewMode(rootId, mode) {
-    const root = document.getElementById(rootId);
-    if (!root) return;
-
-    const base = root.querySelector(".preview-base");
-    const overlay = root.querySelector(".preview-overlay");
-    const blendToggle = root.querySelector(".blend-toggle");
-
-    const enhancedSrc = root.getAttribute("data-enhanced-src");
-    const originalSrc = root.getAttribute("data-original-src");
-
-    if (!base || !overlay || !blendToggle) return;
-
-    if (mode === "enhanced") {
-        base.src = enhancedSrc;
-        overlay.style.display = blendToggle.checked ? "block" : "none";
-        overlay.src = originalSrc;
-    } else if (mode === "original") {
-        base.src = originalSrc;
-        overlay.style.display = "none";
-    }
-}
-
-function toggleBlend(rootId) {
-    const root = document.getElementById(rootId);
-    if (!root) return;
-
-    const overlay = root.querySelector(".preview-overlay");
-    const blendToggle = root.querySelector(".blend-toggle");
-    const opacityWrap = root.querySelector(".opacity-wrap");
-    const base = root.querySelector(".preview-base");
-
-    const originalSrc = root.getAttribute("data-original-src");
-    const enhancedSrc = root.getAttribute("data-enhanced-src");
-
-    if (!overlay || !blendToggle || !opacityWrap || !base) return;
-
-    if (blendToggle.checked) {
-        base.src = enhancedSrc;
-        overlay.src = originalSrc;
-        overlay.style.display = "block";
-        opacityWrap.style.display = "inline-flex";
-    } else {
-        overlay.style.display = "none";
-        opacityWrap.style.display = "none";
-    }
-}
-
-function updateBlendOpacity(rootId, value) {
-    const root = document.getElementById(rootId);
-    if (!root) return;
-
-    const overlay = root.querySelector(".preview-overlay");
-    const valueLabel = root.querySelector(".opacity-value");
-
-    if (!overlay || !valueLabel) return;
-
-    const alpha = Math.max(0, Math.min(100, Number(value))) / 100.0;
-    overlay.style.opacity = alpha.toString();
-    valueLabel.textContent = value + "%";
-}
-</script>
-"""
-
 # ---------------------------------
 # Helpers
 # ---------------------------------
@@ -210,82 +70,105 @@ def get_tile_dimensions(ratio_name: str, tile_preset: str):
 
 def update_tile_display(ratio_name: str, tile_preset: str):
     tile_width, tile_height = get_tile_dimensions(ratio_name, tile_preset)
-    return f"**Tile Width:** {tile_width}px  \n**Tile Height:** {tile_height}px"
+    return (
+        f"**Tile Width:** {tile_width}px  \n"
+        f"**Tile Height:** {tile_height}px"
+    )
 
 
-def save_input_image(input_image_np):
-    img = Image.fromarray(input_image_np).convert("RGB")
-    path = os.path.join(ORIGINAL_DIR, f"{uuid.uuid4().hex}.png")
-    img.save(path, format="PNG", compress_level=0)
-    return img, path
+def format_megapixels(width: int, height: int) -> str:
+    mp = (width * height) / 1_000_000
+    return f"{mp:.2f} MP"
+
+
+def format_file_size(num_bytes: int) -> str:
+    if num_bytes < 1024:
+        return f"{num_bytes} B"
+    if num_bytes < 1024 ** 2:
+        return f"{num_bytes / 1024:.1f} KB"
+    if num_bytes < 1024 ** 3:
+        return f"{num_bytes / (1024 ** 2):.2f} MB"
+    return f"{num_bytes / (1024 ** 3):.2f} GB"
+
+
+def build_stats_markdown(
+    original_width: int,
+    original_height: int,
+    enhanced_width: int,
+    enhanced_height: int,
+    file_size_bytes: int,
+    export_format: str,
+    scale_mode: str,
+):
+    return (
+        f"**Original Dimensions:** {original_width} × {original_height}px  \n"
+        f"**Original Megapixels:** {format_megapixels(original_width, original_height)}  \n\n"
+        f"**Enhanced Dimensions:** {enhanced_width} × {enhanced_height}px  \n"
+        f"**Enhanced Megapixels:** {format_megapixels(enhanced_width, enhanced_height)}  \n\n"
+        f"**Scale Mode:** {scale_mode}  \n"
+        f"**Export Format:** {export_format}  \n"
+        f"**Saved File Size:** {format_file_size(file_size_bytes)}"
+    )
 
 
-def save_output_image(output_img):
+def upscale_once_with_model(
+    img: Image.Image,
+    upscaler,
+    tile_width: int,
+    tile_height: int,
+) -> Image.Image:
+    out = upscaler(
+        img,
+        tiling=True,
+        tile_width=tile_width,
+        tile_height=tile_height,
+    )
+
+    if not isinstance(out, Image.Image):
+        out = Image.fromarray(out)
+
+    return out.convert("RGB")
+
+
+def run_scale_pipeline(
+    img: Image.Image,
+    upscaler,
+    tile_width: int,
+    tile_height: int,
+    scale_mode: str,
+) -> Image.Image:
+    # 4x = one model pass
+    if scale_mode == "4x":
+        return upscale_once_with_model(img, upscaler, tile_width, tile_height)
+
+    # 8x Experimental = one real 4x pass, then high-quality resize to 8x total
+    if scale_mode == "8x Experimental":
+        first = upscale_once_with_model(img, upscaler, tile_width, tile_height)
+        target_w = img.width * 8
+        target_h = img.height * 8
+        return first.resize((target_w, target_h), Image.LANCZOS).convert("RGB")
+
+    # 16x Experimental = two real 4x passes = 16x total
+    if scale_mode == "16x Experimental":
+        first = upscale_once_with_model(img, upscaler, tile_width, tile_height)
+        second = upscale_once_with_model(first, upscaler, tile_width, tile_height)
+        return second.convert("RGB")
+
+    return upscale_once_with_model(img, upscaler, tile_width, tile_height)
+
+
+def save_output_image(output_img: Image.Image, export_format: str):
     output_img = output_img.convert("RGB")
-    path = os.path.join(ENHANCED_DIR, f"{uuid.uuid4().hex}.png")
-    output_img.save(path, format="PNG", compress_level=0)
-    return path
+    file_id = uuid.uuid4().hex
 
+    if export_format == "PNG":
+        path = os.path.join(ENHANCED_DIR, f"{file_id}.png")
+        output_img.save(path, format="PNG", compress_level=0)
+    else:
+        path = os.path.join(ENHANCED_DIR, f"{file_id}.tiff")
+        output_img.save(path, format="TIFF")
 
-def to_file_url(path: str):
-    return f"/gradio_api/file={path}"
-
-
-def build_preview_html(original_path: str, enhanced_path: str):
-    root_id = f"preview_{uuid.uuid4().hex}"
-    original_url = html.escape(to_file_url(original_path), quote=True)
-    enhanced_url = html.escape(to_file_url(enhanced_path), quote=True)
-
-    return f"""
-    <div
-        id="{root_id}"
-        class="preview-shell"
-        data-original-src="{original_url}"
-        data-enhanced-src="{enhanced_url}"
-    >
-        <div class="preview-toolbar">
-            <button class="preview-btn" onclick="setPreviewMode('{root_id}', 'enhanced')">
-                Show Enhanced
-            </button>
-
-            <button class="preview-btn" onclick="setPreviewMode('{root_id}', 'original')">
-                Show Original
-            </button>
-
-            <label class="preview-check">
-                <input
-                    type="checkbox"
-                    class="blend-toggle"
-                    onchange="toggleBlend('{root_id}')"
-                >
-                Blend Overlay Mode
-            </label>
-
-            <div class="preview-slider-wrap opacity-wrap" style="display:none;">
-                <span>Overlay Opacity</span>
-                <input
-                    type="range"
-                    min="0"
-                    max="100"
-                    value="50"
-                    class="opacity-slider"
-                    oninput="updateBlendOpacity('{root_id}', this.value)"
-                >
-                <span class="opacity-value">50%</span>
-            </div>
-        </div>
-
-        <div class="preview-stage">
-            <img class="preview-base" src="{enhanced_url}" alt="Enhanced Preview">
-            <img
-                class="preview-overlay"
-                src="{original_url}"
-                alt="Original Overlay"
-                style="display:none; opacity:0.5;"
-            >
-        </div>
-    </div>
-    """
+    return path
 
 
 # ---------------------------------
@@ -293,38 +176,61 @@ def build_preview_html(original_path: str, enhanced_path: str):
 # ---------------------------------
 
 @spaces.GPU
-def upscale(enhancer_type, ratio_name, tile_preset, input_image):
+def enhance_image(
+    enhancer_type,
+    ratio_name,
+    tile_preset,
+    scale_mode,
+    export_format,
+    input_image,
+):
     if input_image is None:
-        return "<div class='preview-shell'>No preview available yet.</div>", None
+        return (
+            None,
+            None,
+            "No stats available yet."
+        )
 
-    original_img, original_path = save_input_image(input_image)
+    original_img = Image.fromarray(input_image).convert("RGB")
+    original_width, original_height = original_img.size
 
     tile_width, tile_height = get_tile_dimensions(ratio_name, tile_preset)
     upscaler = get_upscaler(enhancer_type)
 
-    enhanced = upscaler(
-        original_img,
-        tiling=True,
+    enhanced_img = run_scale_pipeline(
+        img=original_img,
+        upscaler=upscaler,
         tile_width=tile_width,
         tile_height=tile_height,
+        scale_mode=scale_mode,
     )
 
-    if not isinstance(enhanced, Image.Image):
-        enhanced = Image.fromarray(enhanced)
+    enhanced_width, enhanced_height = enhanced_img.size
+
+    output_path = save_output_image(enhanced_img, export_format)
+    file_size_bytes = os.path.getsize(output_path)
 
-    enhanced_path = save_output_image(enhanced)
-    preview_html = build_preview_html(original_path, enhanced_path)
+    stats_markdown = build_stats_markdown(
+        original_width=original_width,
+        original_height=original_height,
+        enhanced_width=enhanced_width,
+        enhanced_height=enhanced_height,
+        file_size_bytes=file_size_bytes,
+        export_format=export_format,
+        scale_mode=scale_mode,
+    )
 
-    return preview_html, enhanced_path
+    return enhanced_img, output_path, stats_markdown
 
 
 # ---------------------------------
 # UI
 # ---------------------------------
 
-with gr.Blocks(head=HEAD) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# Image Enhancer")
 
+    # 1. Enhancer Type
     with gr.Group():
         enhancer_type = gr.Radio(
             choices=["Anime Enhancer", "Photo Enhancer"],
@@ -332,12 +238,13 @@ with gr.Blocks(head=HEAD) as demo:
             label="Enhancer Type"
         )
 
+    # 2. Combined Tile Settings box
     with gr.Group():
         gr.Markdown("### Tile Settings")
 
         ratio_name = gr.Radio(
             choices=["16:9", "9:16", "4:5", "1:1", "5:4", "2:3", "3:2"],
-            value="16:9",
+            value="1:1",
             label="Aspect Ratio"
         )
 
@@ -348,9 +255,26 @@ with gr.Blocks(head=HEAD) as demo:
         )
 
         tile_display = gr.Markdown(
-            value=update_tile_display("16:9", "768")
+            value=update_tile_display("1:1", "768")
         )
 
+    # Extra processing settings
+    with gr.Group():
+        gr.Markdown("### Output Settings")
+
+        scale_mode = gr.Radio(
+            choices=["4x", "8x Experimental", "16x Experimental"],
+            value="4x",
+            label="Scale Mode"
+        )
+
+        export_format = gr.Radio(
+            choices=["PNG", "TIFF"],
+            value="PNG",
+            label="Export Format"
+        )
+
+    # 3. Input Image
     with gr.Group():
         input_image = gr.Image(
             type="numpy",
@@ -358,20 +282,31 @@ with gr.Blocks(head=HEAD) as demo:
             height=400
         )
 
-    run_button = gr.Button("Enhance 4x")
+    run_button = gr.Button("Enhance Image")
 
+    # 4. Output Preview
     with gr.Group():
         gr.Markdown("### Output Preview")
-        preview_html = gr.HTML(
-            value="<div class='preview-shell'>No preview available yet.</div>"
+        output_preview = gr.Image(
+            type="pil",
+            label="Enhanced Preview",
+            height=400
         )
 
+    # 5. Download box
     with gr.Group():
         gr.Markdown("### New Enhanced Image File")
         download_file = gr.File(
             label="Download new enhanced image file"
         )
 
+    # Stats box
+    with gr.Group():
+        gr.Markdown("### Image Stats")
+        stats_box = gr.Markdown(
+            value="No stats available yet."
+        )
+
     ratio_name.change(
         fn=update_tile_display,
         inputs=[ratio_name, tile_preset],
@@ -385,16 +320,19 @@ with gr.Blocks(head=HEAD) as demo:
     )
 
     run_button.click(
-        fn=upscale,
+        fn=enhance_image,
         inputs=[
             enhancer_type,
             ratio_name,
             tile_preset,
+            scale_mode,
+            export_format,
             input_image,
         ],
         outputs=[
-            preview_html,
+            output_preview,
             download_file,
+            stats_box,
         ],
         show_progress=True
     )