Update app.py

app.py

@@ -1,30 +1,43 @@
-from diffusers import StableDiffusionPipeline
-from PIL import Image
+import gradio as gr
 import torch
+from diffusers import StableDiffusionPipeline
+import subprocess
 
-def generate_image(prompt):
-    """
-    สร้างภาพจากข้อความที่กำหนดโดยใช้ Stable Diffusion.
+# ติดตั้ง flash-attn แม้จะไม่ได้ใช้โดยตรง (ข้าม build CUDA)
+subprocess.run(
+    'pip install flash-attn --no-build-isolation',
+    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+    shell=True
+)
 
-    Args:
-        prompt (str): ข้อความที่ต้องการแปลงเป็นภาพ.
+# ใช้ CPU
+device = "cpu"
 
-    Returns:
-        PIL.Image.Image: รูปภาพที่สร้างขึ้น.
-    """
-    # ตรวจสอบว่ามี CUDA (GPU) หรือไม่ และเลือกอุปกรณ์ที่เหมาะสม
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+# โหลดโมเดลเบา: sd-turbo
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/sd-turbo",
+    torch_dtype=torch.float32
+)
+pipe = pipe.to(device)
+pipe.safety_checker = None  # ปิด safety checker เพื่อความเร็ว
 
-    # โหลดโมเดล Stable Diffusion ที่เทรนไว้ล่วงหน้า
-    model_id = "runwayml/stable-diffusion-v1-5"
-    pipeline = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+# เปิด attention_slicing (แม้บน CPU ก็ช่วยเรื่องหน่วยความจำ)
+pipe.enable_attention_slicing()
 
-    # สร้างภาพ
-    image = pipeline(prompt).images[0]
+# ฟังก์ชันสร้างภาพ
+def generate_image(prompt):
+    result = pipe(prompt, num_inference_steps=10, guidance_scale=3.0)
+    image = result.images[0]
     return image
 
-if __name__ == "__main__":
-    prompt = "A cozy cat sleeping by a fireplace on a snowy evening."
-    generated_image = generate_image(prompt)
-    generated_image.save("cozy_cat.png")
-    print("สร้างภาพเรียบร้อยแล้ว: cozy_cat.png")
+# Gradio UI
+io = gr.Interface(
+    fn=generate_image,
+    inputs=[gr.Textbox(label="Enter your prompt")],
+    outputs=[gr.Image(label="Generated Image")],
+    theme="Yntec/HaleyCH_Theme_Orange",
+    description="⚠ Running on CPU using sd-turbo. Optimized for speed with low inference steps."
+)
+
+# เปิด Gradio ด้วย queue ป้องกันค้างถ้ามีหลายคำสั่ง
+io.queue().launch(debug=True)