import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionInpaintPipeline

def expand_image(image):
    # 画像が256x256より大きい場合、表示用にリサイズ
    max_size = 256  # 最大表示サイズ
    width, height = image.size
    aspect_ratio = width / height  # アスペクト比

    # アスペクト比を維持しながらリサイズ
    if width > max_size or height > max_size:
        if width > height:
            # 横幅が大きい場合、幅をmax_sizeに合わせて高さを自動で調整
            image = image.resize((max_size, int(max_size / aspect_ratio)))
        else:
            # 縦幅が大きい場合、高さをmax_sizeに合わせて幅を自動で調整
            image = image.resize((int(max_size * aspect_ratio), max_size))

    print(f'{image.size=}')  # リサイズ後の画像サイズを確認

    # 画像サイズ（幅）を2.5倍に拡張
    new_width = int(width * 2.5)
    
    # 新しいキャンバスを白で作成し、元画像を中央に配置
    new_image = Image.new("RGB", (new_width, height), (0, 0, 0))  # 黒背景
    new_image.paste(image, ((new_width - width) // 2, 0))  # 元画像を中央に配置

    # マスク画像を作成（変更する部分を白、変更しない部分を黒）
    mask = Image.new("L", (new_width, height), 255)  # 変更する部分を白
    mask.paste(0, ((new_width - width) // 2, 0, (new_width + width) // 2, height))  # 元画像部分を黒 (変更しない部分)

    model_id = "runwayml/stable-diffusion-inpainting"# Stable Diffusion Inpainting パイプラインの準備
    pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float32, use_safetensors=False)
    pipe.to("cpu")  # huggingface は GPUが使えないのでCPUに変更

    prompt = "retain the original image, Extend the current image by adding the most probable realistic elements based on the original content"
    negative_prompt = "repeating, tiling, nsfw, nudity, sexual, violence, explicit, harmful"

    # Inpainting 処理を実行
    output = pipe(prompt=prompt, negative_prompt=negative_prompt, image=new_image, mask_image=mask).images[0]


    return output  # 画像をリサイズ後、結果とともに返す

# Gradioインターフェースの作成
iface = gr.Interface(
    fn=expand_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(),
    title="画像拡張ツール",
    description="入力画像の幅を2.5倍に拡張します。  WIP yet"
)

iface.launch()