Update app.py

app.py

@@ -1,149 +1,63 @@
+import io, os, base64, random, json, tempfile, ipfshttpclient, torch
 import gradio as gr
-import torch
-from diffusers import StableDiffusionPipeline
-import gc
-from PIL import Image, ImageDraw, ImageFont
-import os
+from diffusers import StableDiffusionPipeline, ControlNetModel
+from transformers import pipeline
 
-class ImageGenerator:
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.pipe = None
-        
-    def load_model(self):
-        if self.pipe is None:
-            # 使用較小的模型
-            model_id = "runwayml/stable-diffusion-v1-5"
-            
-            self.pipe = StableDiffusionPipeline.from_pretrained(
-                model_id,
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            
-            if self.device == "cuda":
-                self.pipe = self.pipe.to(self.device)
-                self.pipe.enable_attention_slicing(1)
-                self.pipe.enable_model_cpu_offload()
-    
-    def unload_model(self):
-        if self.pipe is not None:
-            del self.pipe
-            self.pipe = None
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            gc.collect()
-    
-    def generate(self, prompt: str) -> Image.Image:
-        try:
-            self.load_model()
-            
-            image = self.pipe(
-                prompt,
-                num_inference_steps=20,
-                guidance_scale=7.5,
-                height=384,  # 降低解析度以節省記憶體
-                width=384
-            ).images[0]
-            
-            self.unload_model()  # 釋放記憶體
-            return image
-            
-        except Exception as e:
-            error_msg = str(e)
-            print(f"Error: {error_msg}")
-            # 創建錯誤提示圖片
-            error_image = Image.new('RGB', (384, 384), color='white')
-            draw = ImageDraw.Draw(error_image)
-            draw.text((10, 10), f"Error: {error_msg}", fill='black')
-            return error_image
+# ---- 1. 載入四種模型 ----
+sd_pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16
+).to("cuda")
 
-# 全局變數
-generator = ImageGenerator()
+biggan = pipeline("text-to-image", model="osanseviero/BigGAN-deep-128")
 
-def process_prompt(text: str, language: str) -> str:
-    """處理提示詞"""
-    if not text.strip():
-        return "error: empty prompt"
-    
-    if language == "中文":
-        # 添加英文關鍵詞來增強效果
-        enhancers = [
-            "highly detailed",
-            "8k",
-            "high quality",
-            "masterpiece"
-        ]
-        return f"{text}, {', '.join(enhancers)}"
-    return text
+dalle = pipeline("text-to-image", model="dalle-mini/dalle-mini")
 
-def generate_image(prompt: str, language: str) -> Image.Image:
-    """處理圖像生成"""
-    if not prompt.strip():
-        error_image = Image.new('RGB', (384, 384), color='white')
-        draw = ImageDraw.Draw(error_image)
-        draw.text((10, 10), "請輸入提示詞", fill='black')
-        return error_image
-    
-    processed_prompt = process_prompt(prompt, language)
-    return generator.generate(processed_prompt)
+# Hybrid: SD + ControlNet  (邊緣控制範例)
+control = ControlNetModel.from_pretrained(
+    "lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16
+).to("cuda")
 
-# 創建 Gradio 介面
-with gr.Blocks() as demo:
-    gr.Markdown("""
-    ## 🎨 AI 圖像生成系統
-    
-    請在下方輸入您想要生成的圖像描述。
-    
-    注意事項：
-    1. 生成過程可能需要 30-60 秒
-    2. 如果出現錯誤，請稍等片刻後重試
-    3. 中文輸入會自動添加優化關鍵詞
-    """)
-    
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(
-                label="創作提示",
-                placeholder="請輸入描述...",
-                lines=3
-            )
-            language = gr.Radio(
-                choices=["中文", "English"],
-                value="中文",
-                label="語言"
-            )
-            generate_btn = gr.Button("🎨 生成圖像", variant="primary")
-        
-        with gr.Column():
-            output_image = gr.Image(label="生成結果")
-            
-    # 添加示例
-    gr.Examples(
-        examples=[
-            ["日落時分的古老城堡", "中文"],
-            ["peaceful mountain lake", "English"],
-            ["傳統中國水墨畫風格的竹林", "中文"],
-        ],
-        inputs=[prompt, language],
-        outputs=output_image,
-        fn=generate_image,
-        cache_examples=True
-    )
-    
-    # 設置點擊事件
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, language],
-        outputs=output_image
-    )
+def hybrid(prompt):
+    # 先跑一次 SD 取得基底
+    image = sd_pipe(prompt, num_inference_steps=25).images[0]
+    # 取 Canny 邊緣作為條件再跑一次
+    import cv2, numpy as np
+    edges = cv2.Canny(np.array(image), 100, 200)
+    edge_img = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+    final = sd_pipe(prompt, image=edge_img, controlnet=control,
+                    num_inference_steps=20).images[0]
+    return final
 
-# 啟動設定
-if __name__ == "__main__":
-    demo.queue(max_size=1)  # 限制隊列大小
-    demo.launch(
-        show_error=True,
-        share=True,
-        server_name="0.0.0.0"
-    )
+model_dict = {
+    "Diffusion (Stable‑Diffusion‑2‑1)": lambda p: sd_pipe(p, num_inference_steps=30).images[0],
+    "GAN (BigGAN‑128)": lambda p: biggan(p, guidance_scale=5.0)[0]["image"],
+    "Autoregressive (DALL·E Mini)": lambda p: dalle(p, guidance_scale=5.0)[0]["image"],
+    "Hybrid (SD + ControlNet)": hybrid,
+}
+
+# ---- 2. IPFS 上傳 ----
+client = ipfshttpclient.connect("/dns/localhost/tcp/5001/http")  # 或改用 Infura/Pinata
+
+def upload_ipfs(pil_img):
+    buf = io.BytesIO()
+    pil_img.save(buf, format="PNG")
+    buf.seek(0)
+    res = client.add(buf)
+    return f"https://ipfs.io/ipfs/{res['Hash']}"
+
+# ---- 3. Gradio 介面 ----
+def generate(prompt, model_name):
+    img = model_dict[model_name](prompt)
+    cid_url = upload_ipfs(img)
+    return img, cid_url
+
+with gr.Blocks(title="AI Model Mixer") as demo:
+    gr.Markdown("# 🖼️ AI Model Mixer – 4 Models + IPFS")
+    prompt = gr.Textbox(label="輸入提示詞")
+    model = gr.Dropdown(list(model_dict.keys()), value=list(model_dict.keys())[0], label="選擇模型")
+    btn = gr.Button("生成並上傳")
+    gallery = gr.Image(label="輸出影像")
+    cid_out = gr.Textbox(label="IPFS URL")
+    btn.click(generate, inputs=[prompt, model], outputs=[gallery, cid_out])
+
+demo.launch()