Update app.py

app.py

@@ -1,21 +1,28 @@
 import gradio as gr
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 import torch
 from PIL import Image
 
-# 載入 ArchitectureRealMix 模型
 model_id = "stablediffusionapi/architecturerealmix"
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cpu")  # 若無 GPU，可改成 pipe.to("cpu")
+
+# 載入 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):
-    # 調整輸入圖片大小及格式，符合模型需求
-    init_image = input_image.resize((512, 512)).convert("RGB")
-    # 使用文字提示和圖片生成建築模擬圖
-    output = pipe(prompt=prompt, init_image=init_image, strength=strength, guidance_scale=guidance_scale)
-    return output.images[0]
+    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
 
-# 建立 Gradio 使用者介面
 iface = gr.Interface(
     fn=generate_architecture,
     inputs=[
@@ -26,7 +33,7 @@ iface = gr.Interface(
     ],
     outputs=gr.Image(type="pil", label="生成的建築模擬圖"),
     title="ArchitectureRealMix 建築模擬圖生成器",
-    description="利用文字提示搭配圖片上傳生成建築模擬圖。調整圖片強度和引導尺度優化結果。"
+    description="支援以文字描述生成建築模擬圖，亦可上傳圖片結合文字做圖片到圖片生成。"
 )
 
 iface.launch()