Update app.py

app.py

@@ -1,23 +1,54 @@
 import os
-import requests
 from io import BytesIO
-from PIL import Image
+
 import gradio as gr
-from dotenv import load_dotenv
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+
+# =========================
+# DEVICE
+# =========================
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+
+print(f"[INFO] Using device: {DEVICE}, dtype: {DTYPE}")
+
 
 # =========================
-# LOAD ENV
+# MODEL CONFIG
 # =========================
 
-load_dotenv()
+MODEL_OPTIONS = {
+    "Realistic Vision v5.1": "SG161222/Realistic_Vision_V5.1_noVAE",
+    "Stable Diffusion 1.5": "runwayml/stable-diffusion-v1-5",
+    "DreamShaper 8": "Lykon/dreamshaper-8",
+}
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
-HF_MODEL = os.environ.get("HF_MODEL", "playgroundai/playground-v2.5")
+PIPELINES = {}
 
-API_URL = "https://api-inference.huggingface.co/models/stabilityai/sdxl-turbo"
 
-if not HF_TOKEN:
-    print("[WARNING] HF_TOKEN belum di-set di .env")
+def get_pipeline(model_name: str) -> StableDiffusionPipeline:
+    if model_name in PIPELINES:
+        return PIPELINES[model_name]
+
+    repo_id = MODEL_OPTIONS[model_name]
+    print(f"[INFO] Loading model: {model_name} ({repo_id})")
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        repo_id,
+        torch_dtype=DTYPE,
+        safety_checker=None,
+    )
+
+    pipe = pipe.to(DEVICE)
+
+    if DEVICE == "cuda":
+        pipe.enable_xformers_memory_efficient_attention()
+
+    PIPELINES[model_name] = pipe
+    return pipe
 
 
 # =========================
@@ -28,7 +59,7 @@ def auto_prompt(category: str) -> str:
     templates = {
         "Skincare": "Serum skincare botol kaca premium, lighting studio, aesthetic clean look",
         "Makanan/Minuman": "Minuman segar dengan efek splash, lighting vibrant, cocok untuk iklan",
-        "Fashion": "Sepatu fashion modern, lighting studio, katalog e-commerce",
+        "Fashion": "Pakaian atau sepatu fashion modern, lighting studio, katalog e-commerce",
         "Elektronik": "Headphone wireless premium, lighting studio, tampilan high-end",
         "Umum": "Produk premium dengan lighting studio dan background bersih",
     }
@@ -62,41 +93,40 @@ def build_prompt(prompt: str, style: str, category: str, with_model: bool) -> st
         style_map.get(style, ""),
         category_map.get(category, ""),
         model_snippet,
+        "high quality, 4k, detailed",
     ]
 
     return ", ".join([p for p in parts if p])
 
 
 # =========================
-# HUGGINGFACE API CALL
+# GENERATION
 # =========================
 
-def call_huggingface(prompt: str):
-    headers = {
-        "Authorization": f"Bearer {HF_TOKEN}",
-        "Content-Type": "application/json",
-    }
-
-    payload = {
-        "inputs": prompt
-    }
-
-    response = requests.post(API_URL, headers=headers, json=payload)
-    response.raise_for_status()
-
-    img_bytes = response.content
-    img = Image.open(BytesIO(img_bytes))
-
-    return img
+def run(prompt, category, style, with_model, model_choice, steps, guidance, seed):
+    if not prompt or prompt.strip() == "":
+        prompt = auto_prompt(category)
 
+    full_prompt = build_prompt(prompt, style, category, with_model)
 
-# =========================
-# MAIN GENERATION
-# =========================
+    pipe = get_pipeline(model_choice)
+
+    generator = None
+    if seed is not None and seed != "":
+        try:
+            seed_int = int(seed)
+            generator = torch.Generator(device=DEVICE).manual_seed(seed_int)
+        except ValueError:
+            generator = None
+
+    result = pipe(
+        full_prompt,
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        generator=generator,
+    )
 
-def run(prompt, category, style, with_model):
-    full_prompt = build_prompt(prompt, style, category, with_model)
-    img = call_huggingface(full_prompt)
+    img: Image.Image = result.images[0]
     return img
 
 
@@ -104,14 +134,21 @@ def run(prompt, category, style, with_model):
 # GRADIO UI
 # =========================
 
-with gr.Blocks(title="RuangAI – Product Visualizer (Level 2)") as demo:
+with gr.Blocks(title="RuangAI – Product Visualizer (Diffusers)") as demo:
     gr.Markdown("""
-    # 🧴 RuangAI – Product Visualizer (Level 2)
-    Playground v2.5 (HuggingFace Inference API)
+    # 🧴 RuangAI – Product Visualizer (Level 2 – Diffusers Lokal)
+    Tiga model lokal: Realistic Vision v5.1, Stable Diffusion 1.5, DreamShaper 8  
+    **Catatan:** di CPU akan agak lambat, sabar sebentar saat generate 🙏
     """)
 
     with gr.Row():
         with gr.Column():
+            model_choice = gr.Dropdown(
+                list(MODEL_OPTIONS.keys()),
+                value="Realistic Vision v5.1",
+                label="Pilih Model",
+            )
+
             category = gr.Dropdown(
                 ["Umum", "Skincare", "Makanan/Minuman", "Fashion", "Elektronik"],
                 value="Umum",
@@ -135,17 +172,39 @@ with gr.Blocks(title="RuangAI – Product Visualizer (Level 2)") as demo:
                 lines=3,
             )
 
-            auto_btn = gr.Button("Auto Prompt ✨")
-            auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])
+            with gr.Row():
+                auto_btn = gr.Button("Auto Prompt ✨")
+                generate_btn = gr.Button("Generate 🚀")
+
+            steps = gr.Slider(
+                minimum=10,
+                maximum=40,
+                value=25,
+                step=1,
+                label="Inference Steps",
+            )
+
+            guidance = gr.Slider(
+                minimum=3.0,
+                maximum=12.0,
+                value=7.5,
+                step=0.5,
+                label="Guidance Scale",
+            )
 
-            generate_btn = gr.Button("Generate 🚀")
+            seed = gr.Textbox(
+                label="Seed (opsional, untuk hasil konsisten)",
+                placeholder="Kosongkan untuk random",
+            )
 
         with gr.Column():
             output_image = gr.Image(label="Hasil", type="pil")
 
+    auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])
+
     generate_btn.click(
         run,
-        inputs=[prompt, category, style, with_model],
+        inputs=[prompt, category, style, with_model, model_choice, steps, guidance, seed],
         outputs=[output_image],
     )