import gradio as gr
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
import torch
from PIL import Image

model_id = "stablediffusionapi/architecturerealmix"

# 載入 text-to-image 和 image-to-image 兩個管線
pipe_txt2img = StableDiffusionPipeline.from_pretrained(model_id)
pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(model_id)

# 移動到 CPU（根據執行環境可改成 "cuda"）
pipe_txt2img = pipe_txt2img.to("cpu")
pipe_img2img = pipe_img2img.to("cpu")

def generate_architecture(prompt, input_image, strength, guidance_scale):
    if input_image is None:
        # 無上傳圖片，純文字生成
        image = pipe_txt2img(prompt=prompt, guidance_scale=guidance_scale).images[0]
    else:
        # 有圖片，上傳圖片到圖片生成
        init_image = input_image.resize((512, 512)).convert("RGB")
        image = pipe_img2img(prompt=prompt, init_image=init_image, strength=strength, guidance_scale=guidance_scale).images[0]
    return image

iface = gr.Interface(
    fn=generate_architecture,
    inputs=[
        gr.Textbox(lines=2, placeholder="輸入建築物描述文字提示", label="文字提示"),
        gr.Image(type="pil", label="上傳建築設計圖（選擇性）"),
        gr.Slider(0.1, 1.0, value=0.7, step=0.05, label="圖片強度 (strength)"),
        gr.Slider(1.0, 15.0, value=7.5, step=0.1, label="引導尺度 (guidance scale)")
    ],
    outputs=gr.Image(type="pil", label="生成的建築模擬圖"),
    title="ArchitectureRealMix 建築模擬圖生成器",
    description="支援以文字描述生成建築模擬圖，亦可上傳圖片結合文字做圖片到圖片生成。"
)

iface.launch()