Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import numpy as np
 import random
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
+from PIL import Image
 
 # -----------------------------
 # CPU MODE ONLY
@@ -10,35 +11,88 @@ from diffusers import DiffusionPipeline
 device = "cpu"
 torch_dtype = torch.float32
 
-MODEL_ID = "runwayml/stable-diffusion-v1-5"  # model ringan & aman untuk CPU
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
 
-# Load pipeline
-pipe = DiffusionPipeline.from_pretrained(
+# txt2img pipeline
+txt2img_pipe = DiffusionPipeline.from_pretrained(
     MODEL_ID,
     torch_dtype=torch_dtype,
     low_cpu_mem_usage=True,
 )
-pipe.to(device)
+txt2img_pipe.to(device)
+
+# img2img pipeline
+img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch_dtype,
+    low_cpu_mem_usage=True,
+)
+img2img_pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 
 
-def build_prompt(prompt, style):
-    styles = {
+# -----------------------------
+# Prompt builder
+# -----------------------------
+def build_prompt(prompt: str, style: str, category: str) -> str:
+    style_map = {
         "Tanpa gaya": "",
         "Studio": "product photography, clean studio background, soft lighting, high quality",
         "E-commerce": "white background, catalog photo, sharp, high quality, tokopedia, shopee",
         "Pastel": "pastel colors, soft light, aesthetic instagram style",
         "Lifestyle": "realistic lifestyle photography, natural light",
+        "Model Talent": "professional model, commercial photoshoot, studio lighting, natural pose, realistic skin texture, high quality",
+    }
+
+    category_map = {
+        "Umum": "",
+        "Skincare": "skincare product, glossy bottle, premium lighting, beauty aesthetic",
+        "Makanan/Minuman": "food photography, appetizing, vibrant lighting, splash effect",
+        "Fashion": "fashion product, textile detail, clean lighting",
+        "Elektronik": "electronic product, reflective surface, studio lighting",
     }
-    suffix = styles.get(style, "")
-    return f"{prompt}, {suffix}" if suffix else prompt
 
+    s = style_map.get(style, "")
+    c = category_map.get(category, "")
 
-def infer(prompt, negative_prompt, seed, randomize_seed,
-          width, height, guidance_scale, steps,
-          style, num_images):
+    parts = [prompt, s, c]
+    return ", ".join([p for p in parts if p])
 
+
+# -----------------------------
+# Auto prompt generator
+# -----------------------------
+def auto_prompt(category: str) -> str:
+    templates = {
+        "Skincare": "Serum skincare botol kaca premium, tampilan mewah, cocok untuk iklan Instagram",
+        "Makanan/Minuman": "Minuman energi rasa lemon, efek splash, gaya promosi e-commerce",
+        "Fashion": "Sepatu running sport, tampilan katalog, background putih bersih",
+        "Elektronik": "Headphone wireless modern, lighting studio, tampilan premium",
+        "Umum": "Produk premium dengan lighting studio dan background bersih",
+    }
+    return templates.get(category, "Produk premium dengan lighting studio dan background bersih")
+
+
+# -----------------------------
+# Inference
+# -----------------------------
+def generate(
+    mode,
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    steps,
+    style,
+    category,
+    num_images,
+    init_image,
+    strength,
+):
     if not prompt:
         raise gr.Error("Prompt tidak boleh kosong.")
 
@@ -46,58 +100,150 @@ def infer(prompt, negative_prompt, seed, randomize_seed,
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator(device=device).manual_seed(seed)
-    full_prompt = build_prompt(prompt, style)
+    full_prompt = build_prompt(prompt, style, category)
 
     images = []
-    for _ in range(num_images):
-        out = pipe(
-            prompt=full_prompt,
-            negative_prompt=negative_prompt or None,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            generator=generator,
-        )
-        images.append(out.images[0])
+
+    if mode == "Text to Image":
+        for _ in range(num_images):
+            out = txt2img_pipe(
+                prompt=full_prompt,
+                negative_prompt=negative_prompt or None,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=steps,
+                generator=generator,
+            )
+            images.append(out.images[0])
+    else:  # Image to Image
+        if init_image is None:
+            raise gr.Error("Upload gambar produk terlebih dahulu.")
+
+        init_image = init_image.convert("RGB").resize((width, height))
+
+        for _ in range(num_images):
+            out = img2img_pipe(
+                prompt=full_prompt,
+                negative_prompt=negative_prompt or None,
+                image=init_image,
+                strength=strength,
+                guidance_scale=guidance_scale,
+                num_inference_steps=steps,
+                generator=generator,
+            )
+            images.append(out.images[0])
 
     return images, seed
 
 
-with gr.Blocks(title="RuangAI CPU Mode") as demo:
-    gr.Markdown("# 🧴 RuangAI – Product Visualizer (CPU Mode)")
+# -----------------------------
+# UI
+# -----------------------------
+with gr.Blocks(title="RuangAI – Product Visualizer CPU") as demo:
+    gr.Markdown(
+        """
+        # 🧴 RuangAI – Product Visualizer (CPU Mode)
+
+        - **Text to Image**: buat visual produk dari deskripsi  
+        - **Image to Image**: upload foto produk lalu buat versi promosi  
+        - Pilih **gaya visual** dan **kategori produk**  
+        - Gaya **Model Talent** akan menambahkan visualisasi seorang model di hasil gambar
+        """
+    )
+
+    mode = gr.Radio(
+        ["Text to Image", "Image to Image"],
+        value="Text to Image",
+        label="Mode",
+    )
 
     with gr.Row():
-        prompt = gr.Textbox(label="Prompt", placeholder="Deskripsi produk...")
-        run_btn = gr.Button("Generate ✨")
+        category = gr.Dropdown(
+            ["Umum", "Skincare", "Makanan/Minuman", "Fashion", "Elektronik"],
+            value="Umum",
+            label="Kategori Produk",
+        )
+        auto_btn = gr.Button("Auto Prompt ✨")
+
+    prompt = gr.Textbox(
+        label="Prompt",
+        placeholder="Deskripsi produk / ide visual...",
+        lines=3,
+    )
+
+    auto_btn.click(auto_prompt, inputs=[category], outputs=[prompt])
+
+    init_image = gr.Image(
+        label="Upload Gambar (untuk Image to Image)",
+        type="pil",
+    )
 
     with gr.Row():
         style = gr.Dropdown(
-            ["Tanpa gaya", "Studio", "E-commerce", "Pastel", "Lifestyle"],
+            ["Tanpa gaya", "Studio", "E-commerce", "Pastel", "Lifestyle", "Model Talent"],
             value="Studio",
-            label="Gaya visual"
+            label="Gaya visual",
+        )
+        num_images = gr.Slider(
+            1, 4, value=1, step=1, label="Jumlah gambar"
         )
-        num_images = gr.Slider(1, 3, value=1, step=1, label="Jumlah gambar")
 
-    gallery = gr.Gallery(label="Hasil", columns=2, height=512)
+    gallery = gr.Gallery(
+        label="Hasil",
+        columns=2,
+        height=512,
+    )
 
     with gr.Accordion("Advanced Settings", open=False):
-        negative_prompt = gr.Textbox(label="Negative prompt")
-        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
-        randomize_seed = gr.Checkbox(True, label="Randomize seed")
-        width = gr.Slider(256, 768, value=512, step=32, label="Width")
-        height = gr.Slider(256, 768, value=512, step=32, label="Height")
-        guidance_scale = gr.Slider(0, 10, value=7, step=0.5, label="Guidance")
-        steps = gr.Slider(5, 40, value=25, step=1, label="Steps")
+        negative_prompt = gr.Textbox(
+            label="Negative prompt",
+            placeholder="Contoh: blur, low quality, watermark, text, logo",
+        )
+        seed = gr.Slider(
+            0, MAX_SEED, value=0, step=1, label="Seed"
+        )
+        randomize_seed = gr.Checkbox(
+            True, label="Randomize seed"
+        )
+        width = gr.Slider(
+            256, 768, value=512, step=32, label="Width"
+        )
+        height = gr.Slider(
+            256, 768, value=512, step=32, label="Height"
+        )
+        guidance_scale = gr.Slider(
+            0, 10, value=7, step=0.5, label="Guidance"
+        )
+        steps = gr.Slider(
+            5, 40, value=25, step=1, label="Steps"
+        )
+        strength = gr.Slider(
+            0.1, 1.0, value=0.6, step=0.05, label="Strength (img2img)"
+        )
+
+    run_btn = gr.Button("Generate 🚀")
 
     run_btn.click(
-        infer,
+        generate,
         inputs=[
-            prompt, negative_prompt, seed, randomize_seed,
-            width, height, guidance_scale, steps,
-            style, num_images
+            mode,
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            steps,
+            style,
+            category,
+            num_images,
+            init_image,
+            strength,
         ],
-        outputs=[gallery, seed]
+        outputs=[gallery, seed],
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()