app.py

# Copyright 2024 togacat/werecatpete. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import re
import io
import base64
import random
import datetime
from pathlib import Path
from typing import Optional, Dict, List, Union, Tuple

from flask import Flask, request, jsonify, send_from_directory

# ==== ライブラリ分離 ====
try:
    import torch
    from diffusers import DiffusionPipeline
    from PIL import Image, ImageDraw, ImageFont
    from PIL.PngImagePlugin import PngInfo
    HAVE_PIPELINE = True
except Exception:
    torch = None
    DiffusionPipeline = None
    Image = None
    ImageDraw = None
    ImageFont = None
    PngInfo = None
    HAVE_PIPELINE = False

# ==== 基本設定 ====
BASE_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = BASE_DIR / "output"
OUTPUT_DIR.mkdir(exist_ok=True)

# ==== 設定読み込みファイル ====
RANDOM_WORDS_FILE = BASE_DIR / "random_words.txt"

device = "cuda" if (HAVE_PIPELINE and torch.cuda.is_available()) else "cpu"
dtype = torch.float16 if (HAVE_PIPELINE and device == "cuda") else (torch.float32 if HAVE_PIPELINE else None)

if HAVE_PIPELINE:
    print("Loading Mitsua Likes pipeline...")
    try:
        # モデル本体(mitsua-likes)は別途落としておく想定
        pipe = DiffusionPipeline.from_pretrained(
            "mitsua-likes/",
            trust_remote_code=True
        ).to(device, dtype=dtype)
        if hasattr(pipe, "safety_checker"):
            pipe.safety_checker = None
    except Exception as e:
        print("Failed to load Mitsua Likes pipeline:", e)
        pipe = None
else:
    pipe = None

# ==== ランダム単語設定 ====
def load_random_words():
    if RANDOM_WORDS_FILE.exists():
        with open(RANDOM_WORDS_FILE, "r", encoding="utf-8") as f:
            words = [w.strip() for w in f.readlines() if w.strip()]
            if words:
                return words
    return ["night", "fog", "rain", "ruins", "distant city", "reflection"]


RANDOM_WORDS = load_random_words()



def parse_no_logo(value) -> bool:
    """
    JSON から渡ってくる no_logo の値を安全に bool に変換。
    True / "true" / "1" / "yes" / "on" などのみ True とみなし、それ以外は False。
    """
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return bool(value)
    if isinstance(value, str):
        v = value.strip().lower()
        if v in ("1", "true", "yes", "on"):
            return True
        if v in ("0", "false", "no", "off", ""):
            return False
    return False






def parse_disable_character_detection(value) -> bool:
    """JSON から渡ってくる disable_character_detection を安全に bool に変換。"""
    # parse_no_logo と同じ解釈でOK（True/"true"/"1"/"yes"/"on" を True）
    return parse_no_logo(value)


def _has_detected_public_fictional_characters(ret, idx: int = 0) -> bool:
    """
    Mitsua Likes pipeline の戻り値 ret から、公開キャラクター類似が検知されたか判定する。
    ret.detected_public_fictional_characters は [ [str, ...], ... ] を想定。
    """
    try:
        if not hasattr(ret, "detected_public_fictional_characters"):
            return False
        d = getattr(ret, "detected_public_fictional_characters", None)
        if d is None:
            return False
        # list/tuple 想定
        if not isinstance(d, (list, tuple)) or len(d) <= idx:
            return False
        one = d[idx]
        if one is None:
            return False
        if isinstance(one, (list, tuple)) and len(one) > 0:
            return True
        # まれに文字列が入る場合も True 扱い
        if isinstance(one, str) and one.strip():
            return True
        return False
    except Exception:
        return False


def make_character_detection_placeholder(size: Tuple[int, int],
                                         text: str = "detected_public_fictional_characters") -> "Image.Image":
    """真っ白な画像に指定テキストを描画した代替画像を返す。"""
    if Image is None or ImageDraw is None:
        raise RuntimeError("Pillow is not available; cannot create placeholder image.")
    w, h = int(size[0]), int(size[1])
    img = Image.new("RGB", (w, h), (255, 255, 255))
    draw = ImageDraw.Draw(img)

    font = None
    if ImageFont is not None:
        try:
            # Pillow 標準の小さめフォント（環境依存で TTF が無いことがあるため）
            font = ImageFont.load_default()
        except Exception:
            font = None

    # 中央に配置
    try:
        bbox = draw.textbbox((0, 0), text, font=font)
        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
    except Exception:
        try:
            tw, th = draw.textsize(text, font=font)
        except Exception:
            tw, th = (len(text) * 6, 12)

    x = max(0, (w - tw) // 2)
    y = max(0, (h - th) // 2)

    # 白背景なので黒文字
    draw.text((x, y), text, font=font, fill=(0, 0, 0))
    return img
def add_logo_to_image(img: "Image.Image") -> "Image.Image":
    """
    生成画像の下部に "Generated by Mitsua Likes" のロゴテキストを描画して返す。
    - 画像サイズは変更しない（同じキャンバスのまま上書き）
    - できるだけ例外を出さないように実装し、失敗した場合は元画像をそのまま返す
    """
    if Image is None or ImageDraw is None:
        return img

    try:
        # RGBA に変換してオーバーレイレイヤを作成（サイズは元画像と同じ）
        img_rgba = img.convert("RGBA")
        width, height = img_rgba.size

        overlay = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)

        text = "Generated by Mitsua Likes"

        # フォント取得（失敗したら None のまま）
        font = None
        if ImageFont is not None:
            try:
                font = ImageFont.load_default()
            except Exception:
                font = None

        # ロゴ用の帯の高さ（画像の 5% か、最低 24px）
        bar_height = max(height // 20, 24)
        y0 = height - bar_height

        # 半透明の黒帯を下端に描画
        draw.rectangle([0, y0, width, height], fill=(0, 0, 0, 160))

        # テキスト位置（右下寄せ、少し内側）
        padding = 8
        x_text = width - padding
        y_text = y0 + bar_height // 2

        # anchor="rm" で右寄せ＆中央揃え（対応していない Pillow でも無視される）
        try:
            draw.text((x_text, y_text), text, font=font, fill=(255, 255, 255, 255), anchor="rm")
        except TypeError:
            # anchor がサポートされない場合は、右寄せ計算なしで描画
            draw.text((width - padding * 2, y0 + (bar_height - 12) // 2),
                      text, font=font, fill=(255, 255, 255, 255))

        combined = Image.alpha_composite(img_rgba, overlay)
        # 元画像の mode に合わせて戻す（基本は RGB）
        return combined.convert(img.mode)
    except Exception:
        # 何か問題が起きた場合は元画像をそのまま返す
        return img
    try:
        # RGBA に変換してオーバーレイを合成し、その後 RGB に戻す
        img_rgba = img.convert("RGBA")
        overlay = Image.new("RGBA", img_rgba.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)

        # デフォルトフォントを使用
        font = ImageFont.load_default()

        # テキストサイズ取得
        try:
            text_w, text_h = draw.textsize(text, font=font)
        except Exception:
            text_w, text_h = font.getsize(text)

        width, height = img_rgba.size
        padding = 6
        bar_height = text_h + padding * 2
        y0 = height - bar_height

        # 半透明の黒帯
        draw.rectangle([0, y0, width, height], fill=(0, 0, 0, 160))

        # 右下寄せで文字描画
        x_text = width - text_w - padding
        y_text = y0 + (bar_height - text_h) // 2
        draw.text((x_text, y_text), text, font=font, fill=(255, 255, 255, 255))

        combined = Image.alpha_composite(img_rgba, overlay).convert("RGB")
        return combined
    except Exception:
        # 何か問題があればそのまま返す
        return img


def sanitize_for_filename(text: str, max_len: int = 60) -> str:
    text = text.strip()
    text = re.sub(r"\s+", " ", text)
    text = text[:max_len]
    text = re.sub(r'[\\/:*?"<>|]', "_", text)
    if not text:
        text = "prompt"
    return text


def ensure_pipeline_available():
    if pipe is None or torch is None:
        raise RuntimeError(
            "Diffusion pipeline is not available. "
            "Install torch, diffusers, pillow and download 'mitsua-likes' model."
        )


def save_image_with_metadata(img, prompt, negative_prompt,
                             steps, guidance_scale, guidance_rescale,
                             seed, filename_override: Optional[str] = None) -> str:
    if PngInfo is None:
        raise RuntimeError("Pillow is not available; install pillow to save images.")
    metadata = PngInfo()
    param_text = (
        f"{prompt}\n"
        f"Negative prompt: {negative_prompt}\n"
        f"Steps: {steps}, guidance_scale: {guidance_scale}, "
        f"guidance_rescale: {guidance_rescale}, seed: {seed}"
    )
    metadata.add_text("parameters", param_text)

    if filename_override is not None:
        filename = filename_override
    else:
        prompt_part = sanitize_for_filename(prompt)
        now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"{prompt_part}_seed{seed}_{now}.png"

    fullpath = OUTPUT_DIR / filename
    img.save(fullpath, pnginfo=metadata)
    return filename

# ==== latent保存用（i2i用） ====
def save_latent(latent: "torch.Tensor", filename_stem: str) -> str:
    """
    latent を output/<stem>.pt に保存してファイル名を返す。
    """
    if torch is None:
        raise RuntimeError("torch is not available; cannot save latent.")
    path = OUTPUT_DIR / f"{filename_stem}.pt"
    torch.save(latent.detach().cpu(), path)
    return path.name


def parse_parameters_text(param_text: str) -> Optional[Dict]:
    lines = param_text.splitlines()
    if not lines:
        return None

    prompt = lines[0].strip()
    negative_prompt = ""
    steps = 40
    guidance_scale = 6.0
    guidance_rescale = 0.7
    seed = 1000

    for line in lines[1:]:
        if line.startswith("Negative prompt:"):
            negative_prompt = line[len("Negative prompt:"):].strip()
        elif "Steps:" in line:
            parts = [p.strip() for p in line.split(",")]
            for part in parts:
                if part.startswith("Steps:"):
                    steps = int(part.split(":", 1)[1].strip())
                elif part.startswith("guidance_scale:"):
                    guidance_scale = float(part.split(":", 1)[1].strip())
                elif part.startswith("guidance_rescale:"):
                    guidance_rescale = float(part.split(":", 1)[1].strip())
                elif part.startswith("seed:"):
                    seed = int(part.split(":", 1)[1].strip())

    return {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "steps": steps,
        "guidance_scale": guidance_scale,
        "guidance_rescale": guidance_rescale,
        "seed": seed,
    }


def generate_single_image(prompt, negative_prompt,
                          width, height,
                          steps, guidance_scale, guidance_rescale,
                          seed,
                          add_logo: bool = True,
                          enable_character_detection: bool = True):
    """
    通常の 1枚生成。latent はここでは保存しない（inpaint 用に選択されたときだけ再計算して保存する）。
    """
    ensure_pipeline_available()

    generator = torch.Generator().manual_seed(int(seed))

    with torch.no_grad():
        ret = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=float(guidance_scale),
            guidance_rescale=float(guidance_rescale),
            generator=generator,
            width=int(width),
            height=int(height),
            num_inference_steps=int(steps),
            # return_latents はデフォルト False
        )

    try:
        if hasattr(ret, "detected_public_fictional_characters"):
            print("Similarity Restriction:", ret.detected_public_fictional_characters[0])
        if hasattr(ret, "detected_public_fictional_characters_info"):
            print("Similarity Measure:")
            for k, v in ret.detected_public_fictional_characters_info[0].items():
                print(f"{k} : {v:.3%}")
    except Exception:
        pass

    img = ret.images[0]
    detected = bool(enable_character_detection and _has_detected_public_fictional_characters(ret, 0))
    if detected:
        img = make_character_detection_placeholder(img.size, "detected_public_fictional_characters")
    elif add_logo:
        img = add_logo_to_image(img)
    filename = save_image_with_metadata(
        img, prompt, negative_prompt, steps, guidance_scale, guidance_rescale, seed
    )

    return {
        "filename": filename,
        "url": f"/output/{filename}",
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "steps": steps,
        "guidance_scale": guidance_scale,
        "guidance_rescale": guidance_rescale,
        "seed": seed,
        "latent_filename": None,  # latent は生成していない
    }


def generate_grid9(prompt, negative_prompt,
                   width, height,
                   steps, guidance_scale, guidance_rescale,
                   seeds,
                   add_logo: bool = True,
                   enable_character_detection: bool = True):
    """
    3x3 (9枚) をまとめて生成。ここでも latent は保存しない。
    """
    ensure_pipeline_available()

    if len(seeds) != 9:
        raise ValueError("seeds must have length 9")

    prompts = [prompt] * 9
    negative_prompts = [negative_prompt] * 9
    generators = [torch.Generator().manual_seed(int(s)) for s in seeds]

    with torch.no_grad():
        ret = pipe(
            prompt=prompts,
            negative_prompt=negative_prompts,
            guidance_scale=float(guidance_scale),
            guidance_rescale=float(guidance_rescale),
            generator=generators,
            width=int(width),
            height=int(height),
            num_inference_steps=int(steps),
        )

    images_info = []
    imgs = ret.images

    for idx, (seed, img) in enumerate(zip(seeds, imgs)):
        detected = bool(enable_character_detection and _has_detected_public_fictional_characters(ret, idx))
        if detected:
            img_to_save = make_character_detection_placeholder(img.size, "detected_public_fictional_characters")
        elif add_logo:
            img_to_save = add_logo_to_image(img)
        else:
            img_to_save = img
        filename = save_image_with_metadata(
            img_to_save, prompt, negative_prompt, steps, guidance_scale, guidance_rescale, seed
        )
        images_info.append({
            "filename": filename,
            "url": f"/output/{filename}",
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "steps": steps,
            "guidance_scale": guidance_scale,
            "guidance_rescale": guidance_rescale,
            "seed": seed,
            "latent_filename": None,
        })

    return images_info


def get_latent_path_for_image(filename: str) -> Optional[Path]:
    stem = Path(filename).stem
    pt = OUTPUT_DIR / f"{stem}.pt"
    if pt.exists():
        return pt
    return None


def reconstruct_latent_for_image(filename: str) -> Optional[Path]:
    """
    PNG の parameters から latent を再計算して保存し、その Path を返す。
    失敗したら None を返す。
    """
    if Image is None or torch is None or pipe is None:
        return None

    png_path = OUTPUT_DIR / filename
    if not png_path.exists():
        return None

    img = Image.open(png_path)
    width, height = img.width, img.height
    param_text = img.info.get("parameters")
    if not param_text:
        return None

    meta = parse_parameters_text(param_text)
    if meta is None:
        return None

    prompt = meta["prompt"]
    negative_prompt = meta["negative_prompt"]
    steps = meta["steps"]
    guidance_scale = meta["guidance_scale"]
    guidance_rescale = meta["guidance_rescale"]
    seed = meta["seed"]

    try:
        ensure_pipeline_available()
        generator = torch.Generator().manual_seed(int(seed))
        with torch.no_grad():
            ret = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=float(guidance_scale),
                guidance_rescale=float(guidance_rescale),
                generator=generator,
                width=int(width),
                height=int(height),
                num_inference_steps=int(steps),
                return_latents=True,
            )
        latents = getattr(ret, "latents", None)
        if isinstance(latents, torch.Tensor):
            latents_1 = latents[0]
        elif isinstance(latents, (list, tuple)) and len(latents) > 0:
            latents_1 = latents[0]
        else:
            return None

        stem = Path(filename).stem
        latent_name = save_latent(latents_1, stem)
        return OUTPUT_DIR / latent_name
    except Exception:
        # ここで詳細なログを入れたければ print(traceback.format_exc()) など
        return None


def decode_mask_from_base64(png_base64: str, target_latent_shape):
    """
    dataURL (image/png;base64,...) を受け取り、
    latent の (H, W) に縮小した 0〜1 の torch.Tensor mask を返す。
    """
    if Image is None or torch is None:
        raise RuntimeError("Pillow/torch are required for mask processing.")

    if png_base64.startswith("data:"):
        png_base64 = png_base64.split(",", 1)[1]

    raw = base64.b64decode(png_base64)
    img = Image.open(io.BytesIO(raw)).convert("RGBA")
    # alpha を利用 (0〜255)
    alpha = img.split()[-1]  # A
    # 閾値10で二値化
    alpha = alpha.point(lambda v: 255 if v > 10 else 0)

    # latent 解像度へ縮小
    _, _, h_latent, w_latent = target_latent_shape
    alpha_small = alpha.resize((w_latent, h_latent), resample=Image.NEAREST)

    # 0〜1 の Tensor
    import numpy as np
    alpha_np = np.array(alpha_small, dtype="float32") / 255.0
    mask = torch.from_numpy(alpha_np)
    mask = mask.clamp(0.0, 1.0)
    mask = mask.unsqueeze(0).unsqueeze(0)  # (1,1,h,w)
    return mask.to(device=device, dtype=dtype)


app = Flask(__name__, static_folder="static", static_url_path="/static")


@app.route("/")
def index():
    return send_from_directory(app.static_folder, "index.html")


@app.route("/output/<path:filename>")
def serve_output(filename):
    return send_from_directory(OUTPUT_DIR, filename)


@app.route("/api/generate-one", methods=["POST"])
def api_generate_one():
    data = request.get_json(force=True)
    prompt = data.get("prompt", "").strip()
    negative_prompt = data.get("negative_prompt", "").strip()
    width = int(data.get("width", 672))
    height = int(data.get("height", 896))
    steps = int(data.get("steps", 40))
    guidance_scale = float(data.get("guidance_scale", 6.0))
    guidance_rescale = float(data.get("guidance_rescale", 0.7))

    seed_param = data.get("seed", None)
    if seed_param is None or str(seed_param) == "":
        seed = random.randint(0, 2**31 - 1)
    else:
        seed = int(seed_param)

    no_logo = parse_no_logo(data.get("no_logo", False))
    disable_character_detection = parse_disable_character_detection(data.get("disable_character_detection", False))


    try:
        info = generate_single_image(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=steps,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            seed=seed,
            add_logo=not no_logo,
            enable_character_detection=not disable_character_detection,
        )
        return jsonify(info)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/api/generate-grid9", methods=["POST"])
def api_generate_grid9():
    data = request.get_json(force=True)
    prompt = data.get("prompt", "").strip()
    negative_prompt = data.get("negative_prompt", "").strip()
    width = int(data.get("width", 672))
    height = int(data.get("height", 896))
    steps = int(data.get("steps", 40))
    guidance_scale = float(data.get("guidance_scale", 6.0))
    guidance_rescale = float(data.get("guidance_rescale", 0.7))

    first_seed_param = data.get("first_seed", None)

    seeds = []
    if first_seed_param is not None and str(first_seed_param) != "":
        seed0 = int(first_seed_param)
    else:
        seed0 = random.randint(0, 2**31 - 1)
    seeds.append(seed0)
    for _ in range(8):
        seeds.append(random.randint(0, 2**31 - 1))

    no_logo = parse_no_logo(data.get("no_logo", False))

    disable_character_detection = parse_disable_character_detection(data.get("disable_character_detection", False))

    try:
        images_info = generate_grid9(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=steps,
            guidance_scale=guidance_scale,
            guidance_rescale=guidance_rescale,
            seeds=seeds,
            add_logo=not no_logo,
            enable_character_detection=not disable_character_detection,
        )
        return jsonify({"images": images_info})
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route("/api/random-word", methods=["GET"])
def api_random_word():
    word = random.choice(RANDOM_WORDS)
    return jsonify({"word": word})


@app.route("/api/upload-image", methods=["POST"])
def api_upload_image():
    if Image is None:
        return jsonify({"error": "Pillow is not available; install pillow to use this feature."}), 500

    file = request.files.get("file")
    if not file:
        return jsonify({"error": "no file"}), 400

    original_name = sanitize_for_filename(file.filename or "uploaded")
    now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = f"uploaded_{now}_{original_name}.png"
    save_path = OUTPUT_DIR / filename
    file.save(save_path)

    img = Image.open(save_path)
    param_text = img.info.get("parameters")
    if not param_text:
        meta = {
            "prompt": "",
            "negative_prompt": "",
            "steps": 40,
            "guidance_scale": 6.0,
            "guidance_rescale": 0.7,
            "seed": random.randint(0, 2**31 - 1),
        }
    else:
        parsed = parse_parameters_text(param_text)
        if parsed is None:
            meta = {
                "prompt": "",
                "negative_prompt": "",
                "steps": 40,
                "guidance_scale": 6.0,
                "guidance_rescale": 0.7,
                "seed": random.randint(0, 2**31 - 1),
            }
        else:
            meta = parsed

    meta.update({
        "filename": filename,
        "url": f"/output/{filename}",
        "width": img.width,
        "height": img.height,
        "latent_filename": None,
    })

    return jsonify(meta)


@app.route("/api/inpaint-init", methods=["POST"])
def api_inpaint_init():
    """
    inpaint モード開始時に呼び出す。
    - 既に latent .pt があればその名前を返す。
    - 無ければ、PNG の parameters から同じ条件で再生成して latent だけ保存する。
    """
    data = request.get_json(force=True)
    filename = data.get("filename")
    if not filename:
        return jsonify({"error": "filename is required"}), 400

    # すでに .pt があるか？
    latent_path = get_latent_path_for_image(filename)
    if latent_path is not None:
        return jsonify({"latent_filename": latent_path.name})

    # 無ければ PNG からメタ情報を読み出して latent を再計算
    if Image is None or torch is None or pipe is None:
        return jsonify({"error": "pipeline not available for latent reconstruction"}), 500

    png_path = OUTPUT_DIR / filename
    if not png_path.exists():
        return jsonify({"error": f"image not found: {filename}"}), 400

    img = Image.open(png_path)
    width, height = img.width, img.height
    param_text = img.info.get("parameters")
    if not param_text:
        return jsonify({"error": "no parameters in PNG; cannot reconstruct latent"}), 400

    meta = parse_parameters_text(param_text)
    if meta is None:
        return jsonify({"error": "failed to parse parameters; cannot reconstruct latent"}), 400

    prompt = meta["prompt"]
    negative_prompt = meta["negative_prompt"]
    steps = meta["steps"]
    guidance_scale = meta["guidance_scale"]
    guidance_rescale = meta["guidance_rescale"]
    seed = meta["seed"]

    try:
        ensure_pipeline_available()
        generator = torch.Generator().manual_seed(int(seed))
        with torch.no_grad():
            ret = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=float(guidance_scale),
                guidance_rescale=float(guidance_rescale),
                generator=generator,
                width=int(width),
                height=int(height),
                num_inference_steps=int(steps),
                return_latents=True,
            )
        latents = getattr(ret, "latents", None)
        if isinstance(latents, torch.Tensor):
            latents_1 = latents[0]
        elif isinstance(latents, (list, tuple)) and len(latents) > 0:
            latents_1 = latents[0]
        else:
            return jsonify({"error": "pipeline did not return latents"}), 500

        stem = Path(filename).stem
        latent_filename = save_latent(latents_1, stem)
        return jsonify({"latent_filename": latent_filename})
    except Exception as e:
        return jsonify({"error": f"latent reconstruction failed: {e}"}), 500


def get_unique_inpaint_names(original_filename: str, original_latent_filename: str) -> (str, str):
    """
    もとの画像/latent のファイル名から、上書きしない I<n>_ プレフィックス付きの
    画像ファイル名と latent ステムを返す。
    例: original.png -> I1_original.png, I1_original_latentstem
    """
    base_name = Path(original_filename).name
    latent_stem = Path(original_latent_filename).stem if original_latent_filename else Path(original_filename).stem

    # 既存の I<n>_base_name を探して、次の番号を振る
    n = 1
    while True:
        candidate_img = f"I{n}_" + base_name
        if not (OUTPUT_DIR / candidate_img).exists():
            break
        n += 1

    new_image_name = f"I{n}_" + base_name
    new_latent_stem = f"I{n}_" + latent_stem
    return new_image_name, new_latent_stem


@app.route("/api/inpaint-repaint", methods=["POST"])
def api_inpaint_repaint():
    """
    latent + mask + prompt から inpaint して、新しい画像＆latent を保存。
    Mitsua Likes の pipeline 側に inpaint_internal_latents(latents, mask, ...) を
    追加してある前提。
    """
    if pipe is None or torch is None:
        return jsonify({"error": "pipeline not available"}), 500
    if not hasattr(pipe, "inpaint_internal_latents"):
        return jsonify({
            "error": "pipe.inpaint_internal_latents が実装されていません。"
                     "Mitsua Likes の pipeline に inpaint_internal_latents を追加してください。"
        }), 500

    data = request.get_json(force=True)
    filename = data.get("filename")
    latent_filename = data.get("latent_filename")
    prompt = data.get("prompt", "").strip()
    negative_prompt = data.get("negative_prompt", "").strip()
    width = int(data.get("width", 672))
    height = int(data.get("height", 896))
    steps = int(data.get("steps", 40))
    guidance_scale = float(data.get("guidance_scale", 6.0))
    guidance_rescale = float(data.get("guidance_rescale", 0.7))
    no_logo = parse_no_logo(data.get("no_logo", False))
    disable_character_detection = parse_disable_character_detection(data.get("disable_character_detection", False))
    seed = int(data.get("seed", 0))  # seed はメタ情報用
    strength = float(data.get("strength", 0.5))
    mask_png = data.get("mask_png")

    # 追加: 画像全体を i2i するモードかどうか
    i2i_mode = bool(data.get("i2i_mode", False))
    i2i_mask_pct = data.get("i2i_mask_pct", None)

    if not filename:
        return jsonify({"error": "filename is required"}), 400
    if not mask_png and not i2i_mode:
        return jsonify({"error": "mask_png is required"}), 400


    # 1) まずフロントからの latent_filename を試す
    latent_path: Optional[Path] = None
    if latent_filename:
        candidate = OUTPUT_DIR / latent_filename
        if candidate.exists():
            latent_path = candidate

    # 2) 無い場合は、同名 PNG 用の既存 .pt を探す
    if latent_path is None:
        lp = get_latent_path_for_image(filename)
        if lp is not None and lp.exists():
            latent_path = lp
            latent_filename = lp.name

    # 3) それでも無ければ、PNG から latent を再構成して保存する
    if latent_path is None:
        lp = reconstruct_latent_for_image(filename)
        if lp is None or not lp.exists():
            # ここで初めて「latent が無い」とエラーにする
            return jsonify({
                "error": "latent (.pt) が見つかりません（latents 付きで生成された画像か、PNG+PT をドロップしてください）"
            }), 400
        latent_path = lp
        latent_filename = lp.name

    # latent 読み込み
    latents = torch.load(latent_path, map_location=device)
    if latents.dim() == 3:
        latents = latents.unsqueeze(0)  # (C,H,W) -> (1,C,H,W)

    # mask を latent 解像度に変換
    if i2i_mode:
        # 画像全体に一様マスク (0.0〜1.0)
        pct = 100.0
        if i2i_mask_pct is not None:
            try:
                pct = float(i2i_mask_pct)
            except (TypeError, ValueError):
                pct = 100.0
        pct = max(0.0, min(100.0, pct))
        mask_value = pct / 100.0

        # latents.shape: (B, C, H, W)
        _, _, h, w = latents.shape
        # 1バッチ前提。必要なら B に合わせて拡張可能
        mask_tensor = torch.full((1, 1, h, w), mask_value, dtype=dtype)
    else:
        # 通常の手描きマスク
        mask_tensor = decode_mask_from_base64(mask_png, latents.shape)

    # inpaint 実行
    try:
        with torch.no_grad():
            ret = pipe.inpaint_internal_latents(
                latents=latents.to(device=device, dtype=dtype),
                mask=mask_tensor,
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                guidance_rescale=guidance_rescale,
                output_type="pil",
            )
    except Exception as e:
        return jsonify({"error": f"inpaint_internal_latents failed: {e}"}), 500

    # 新しい画像＆latent
    img = ret.images[0]
    detected = bool((not disable_character_detection) and _has_detected_public_fictional_characters(ret, 0))
    if detected:
        img = make_character_detection_placeholder(img.size, "detected_public_fictional_characters")
    elif not no_logo:
        img = add_logo_to_image(img)
    new_latents = getattr(ret, "latents", None)
    if isinstance(new_latents, torch.Tensor) and new_latents.dim() == 4:
        new_latent_1 = new_latents[0]
    elif isinstance(new_latents, (list, tuple)) and len(new_latents) > 0:
        new_latent_1 = new_latents[0]
    else:
        new_latent_1 = None

    # ファイル名: もとのファイル名の先頭に I<n>_ を付けてナンバリングし、上書きを避ける
    new_image_name, new_latent_name_stem = get_unique_inpaint_names(filename, latent_filename)

    # 画像保存（メタ情報付き）
    saved_image_name = save_image_with_metadata(
        img, prompt, negative_prompt, steps, guidance_scale, guidance_rescale, seed,
        filename_override=new_image_name
    )

    # latent 保存
    if new_latent_1 is not None:
        latent_saved_name = save_latent(new_latent_1, new_latent_name_stem)
    else:
        latent_saved_name = None

    info = {
        "filename": saved_image_name,
        "url": f"/output/{saved_image_name}",
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "steps": steps,
        "guidance_scale": guidance_scale,
        "guidance_rescale": guidance_rescale,
        "seed": seed,
        "latent_filename": latent_saved_name,
    }
    return jsonify(info)

@app.route("/api/inpaint-upload", methods=["POST"])
def api_inpaint_upload():
    """
    inpaint モード用のアップロード。
    PNG と対応する .pt をまとめてドラッグ＆ドロップした場合を想定。
    - PNG からメタ情報を読み取り
    - .pt は output/ にコピー
    """
    if Image is None:
        return jsonify({"error": "Pillow is not available"}), 500

    files = request.files.getlist("files")
    if not files:
        return jsonify({"error": "no files"}), 400

    png_file = None
    pt_file = None
    for f in files:
        name = f.filename or ""
        lower = name.lower()
        if lower.endswith(".png"):
            png_file = f
        elif lower.endswith(".pt"):
            pt_file = f

    if png_file is None:
        return jsonify({"error": "png file is required"}), 400

    # PNG を保存
    original_name = sanitize_for_filename(png_file.filename or "uploaded")
    now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = f"inpaint_uploaded_{now}_{original_name}.png"
    save_path = OUTPUT_DIR / filename
    png_file.save(save_path)

    img = Image.open(save_path)
    param_text = img.info.get("parameters")
    if not param_text:
        meta = {
            "prompt": "",
            "negative_prompt": "",
            "steps": 40,
            "guidance_scale": 6.0,
            "guidance_rescale": 0.7,
            "seed": random.randint(0, 2**31 - 1),
        }
    else:
        parsed = parse_parameters_text(param_text)
        if parsed is None:
            meta = {
                "prompt": "",
                "negative_prompt": "",
                "steps": 40,
                "guidance_scale": 6.0,
                "guidance_rescale": 0.7,
                "seed": random.randint(0, 2**31 - 1),
            }
        else:
            meta = parsed

    latent_filename = None
    if pt_file is not None:
        latent_name = f"{Path(filename).stem}.pt"
        latent_path = OUTPUT_DIR / latent_name
        pt_file.save(latent_path)
        latent_filename = latent_name

    meta.update({
        "filename": filename,
        "url": f"/output/{filename}",
        "width": img.width,
        "height": img.height,
        "latent_filename": latent_filename,
    })

    return jsonify(meta)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)