app.py

import os
import gc
import uuid
import time
import random
import requests

import torch
import gradio as gr

from PIL import Image
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline


# =========================
# DEVICE
# =========================
device = "cpu"

# =========================
# DREAMSHAPER MODEL CONFIG
# =========================
MODEL_ID = "Lykon/dreamshaper-8"

# =========================
# GROQ API CONFIG
# =========================
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"

# =========================
# GLOBAL STATE
# =========================
current_pipe = None
current_mode = None
history = []


# =========================
# STYLE PRESETS
# =========================
STYLE_PRESETS = {
    "None": "",
    "Realistic": "photorealistic, realistic lighting, realistic details, natural skin texture",
    "Anime": "anime style, vibrant colors, clean line art, detailed anime illustration",
    "Cinematic": "cinematic lighting, dramatic shadows, film still, movie composition",
    "Fantasy": "fantasy art, magical atmosphere, epic composition, highly detailed",
    "Product": "professional product photography, studio lighting, clean background",
}


# =========================
# NEGATIVE PROMPT
# =========================
NEGATIVE_PROMPT = (
    "worst quality, low quality, blurry, bad anatomy, bad proportions, "
    "extra limbs, missing limbs, distorted face, deformed hands, bad hands, "
    "extra fingers, missing fingers, duplicated body parts, ugly, watermark, "
    "text, logo, signature, jpeg artifacts, oversaturated"
)


# =========================
# SIMPLE LOCAL TRANSLATOR ID TO EN
# FALLBACK JIKA GROQ API TIDAK AKTIF
# =========================
def translate_prompt(prompt):
    if not prompt:
        return ""

    mapping = {
        "pria": "man",
        "laki-laki": "man",
        "wanita": "woman",
        "perempuan": "woman",
        "anak": "child",
        "wajah": "face",
        "potret": "portrait",
        "gunung": "mountain",
        "pantai": "beach",
        "kota": "city",
        "malam": "night",
        "siang": "day",
        "pagi": "morning",
        "sore": "afternoon",
        "senja": "sunset",
        "laut": "sea",
        "langit": "sky",
        "hutan": "forest",
        "mobil": "car",
        "motor": "motorcycle",
        "rumah": "house",
        "jalan": "street",
        "sungai": "river",
        "danau": "lake",
        "kucing": "cat",
        "anjing": "dog",
        "burung": "bird",
        "bunga": "flower",
        "pohon": "tree",
        "robot": "robot",
        "hujan": "rain",
        "salju": "snow",
        "api": "fire",
        "air": "water",
        "cantik": "beautiful",
        "tampan": "handsome",
        "futuristik": "futuristic",
        "tradisional": "traditional",
        "jawa": "javanese",
        "indonesia": "indonesia",
        "desa": "village",
        "sawah": "rice field",
        "kerajaan": "kingdom",
        "istana": "palace",
        "emas": "gold",
        "perak": "silver",
        "hitam": "black",
        "putih": "white",
        "merah": "red",
        "biru": "blue",
        "hijau": "green",
        "kuning": "yellow",
    }

    prompt = prompt.lower()

    for k, v in mapping.items():
        prompt = prompt.replace(k, v)

    return prompt


# =========================
# DETECT PROMPT TYPE
# =========================
def detect_type(prompt):
    p = prompt.lower()

    portrait_keywords = [
        "person", "man", "woman", "girl", "boy",
        "face", "portrait", "child", "people"
    ]

    scene_keywords = [
        "mountain", "city", "forest", "beach", "sky",
        "sea", "street", "river", "lake", "tree",
        "village", "landscape", "room", "house",
        "rice field", "palace"
    ]

    if any(k in p for k in portrait_keywords):
        return "portrait"

    if any(k in p for k in scene_keywords):
        return "scene"

    return "object"


# =========================
# LOCAL PROMPT ENGINE
# =========================
def build_prompt_local(prompt, style):
    """
    Fallback prompt engine lokal jika Groq API tidak tersedia.
    """
    prompt = translate_prompt(prompt)
    style_text = STYLE_PRESETS.get(style, "")

    base = "masterpiece, best quality, ultra detailed, sharp focus, high detail"
    detail = "intricate details, detailed texture, balanced composition"

    ptype = detect_type(prompt)

    if ptype == "portrait":
        extra = (
            "detailed face, natural skin texture, expressive eyes, "
            "soft lighting, portrait composition, depth of field, "
            "realistic facial proportions"
        )
    elif ptype == "scene":
        extra = (
            "wide shot, cinematic composition, environmental detail, "
            "atmospheric lighting, realistic perspective, immersive scene"
        )
    else:
        extra = (
            "centered composition, studio lighting, clean background, "
            "sharp object details, professional photography"
        )

    parts = [
        prompt,
        style_text,
        base,
        detail,
        extra
    ]

    return ", ".join([p for p in parts if p])


# =========================
# GROQ PROMPT ENHANCER
# =========================
def enhance_prompt_with_groq(user_prompt, style):
    """
    Mengubah prompt Bahasa Indonesia menjadi prompt Bahasa Inggris
    yang lebih optimal untuk DreamShaper / Stable Diffusion menggunakan Groq API.
    """

    if not user_prompt:
        return ""

    if not GROQ_API_KEY:
        print("GROQ_API_KEY belum diset. Menggunakan prompt lokal.")
        return build_prompt_local(user_prompt, style)

    system_prompt = """
You are an expert prompt engineer for Stable Diffusion DreamShaper.

Your task:
- Convert Indonesian image prompts into high-quality English prompts.
- Preserve the user's exact visual intent.
- Do not change the main subject.
- Do not add unrelated objects.
- Add useful visual details only when helpful:
  composition, lighting, atmosphere, camera angle, lens, texture, detail level, realism, style.
- Make the final prompt suitable for Stable Diffusion / DreamShaper.
- Return only one final English prompt.
- Do not use markdown.
- Do not use bullet points.
- Do not explain anything.

Important:
- Keep the output concise but rich in visual detail.
- Avoid overly long prompts.
- Keep the prompt general-audience safe.
- If the request is unsafe or inappropriate, rewrite it into a safe, non-explicit, general-audience visual prompt.
"""

    selected_style = STYLE_PRESETS.get(style, "")

    user_message = f"""
User Indonesian prompt:
{user_prompt}

Selected style:
{style}

Style keyword:
{selected_style}

Create one optimized English prompt for DreamShaper.
"""

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": user_message
            }
        ],
        "temperature": 0.35,
        "max_tokens": 300
    }

    try:
        response = requests.post(
            GROQ_API_URL,
            headers=headers,
            json=payload,
            timeout=30
        )

        response.raise_for_status()

        data = response.json()

        enhanced_prompt = data["choices"][0]["message"]["content"].strip()
        enhanced_prompt = enhanced_prompt.replace("\n", " ").strip()

        if not enhanced_prompt:
            return build_prompt_local(user_prompt, style)

        return enhanced_prompt

    except Exception as e:
        print("Groq prompt enhancer error:", e)
        print("Fallback ke prompt lokal.")
        return build_prompt_local(user_prompt, style)


# =========================
# MAIN PROMPT BUILDER
# =========================
def build_prompt(prompt, style, use_groq=True):
    """
    Prompt utama.
    Jika use_groq aktif dan API key tersedia, prompt diproses Groq.
    Jika gagal, fallback ke prompt lokal.
    """
    if use_groq:
        return enhance_prompt_with_groq(prompt, style)

    return build_prompt_local(prompt, style)


# =========================
# LOAD PIPELINE DREAMSHAPER ONLY
# =========================
def load_pipe(mode):
    """
    mode:
    - txt2img
    - img2img
    """
    global current_pipe, current_mode

    if current_pipe is not None and current_mode == mode:
        return current_pipe

    if current_pipe is not None:
        print("Clearing previous pipeline...")
        del current_pipe
        current_pipe = None
        current_mode = None
        gc.collect()

    print(f"Loading DreamShaper pipeline mode: {mode}")

    if mode == "txt2img":
        pipe = StableDiffusionPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True
        )

    elif mode == "img2img":
        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True
        )

    else:
        raise ValueError("Unknown pipeline mode")

    pipe.to(device)
    pipe.enable_attention_slicing()

    # Sesuai kode awal Anda
    pipe.safety_checker = None

    current_pipe = pipe
    current_mode = mode

    return pipe


# =========================
# UTILITIES
# =========================
def normalize_seed(seed):
    try:
        seed = int(seed)
    except Exception:
        seed = -1

    if seed == -1:
        seed = random.randint(0, 999999)

    return seed


def normalize_size(value, default=512):
    try:
        value = int(value)
    except Exception:
        value = default

    value = max(256, min(value, 768))

    # Stable Diffusion lebih aman kelipatan 8
    value = value - (value % 8)

    return value


def upscale(img):
    return img.resize(
        (img.width * 2, img.height * 2),
        Image.LANCZOS
    )


def save_image(img):
    os.makedirs("outputs", exist_ok=True)

    filename = f"{uuid.uuid4().hex}.png"
    path = os.path.join("outputs", filename)

    img.save(path)

    return path


def prepare_img2img_input(input_image):
    if input_image is None:
        return None

    if not isinstance(input_image, Image.Image):
        input_image = Image.fromarray(input_image)

    input_image = input_image.convert("RGB")

    input_image.thumbnail(
        (768, 768),
        Image.LANCZOS
    )

    w = input_image.width - (input_image.width % 8)
    h = input_image.height - (input_image.height % 8)

    w = max(256, w)
    h = max(256, h)

    input_image = input_image.resize(
        (w, h),
        Image.LANCZOS
    )

    return input_image


# =========================
# TEXT TO IMAGE GENERATOR
# =========================
def generate_text_to_image(
    prompt,
    style,
    use_groq,
    steps,
    width,
    height,
    seed,
    up
):
    pipe = load_pipe("txt2img")

    final_prompt = build_prompt(prompt, style, use_groq)

    print("FINAL TXT2IMG PROMPT:", final_prompt)

    width = normalize_size(width)
    height = normalize_size(height)
    seed = normalize_seed(seed)

    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=final_prompt,
        negative_prompt=NEGATIVE_PROMPT,
        width=width,
        height=height,
        guidance_scale=7.5,
        num_inference_steps=int(steps),
        generator=generator
    ).images[0]

    if up:
        image = upscale(image)

    history.append(image)

    file_path = save_image(image)

    return image, file_path, seed, final_prompt, history


# =========================
# IMAGE TO IMAGE GENERATOR
# =========================
def generate_image_to_image(
    prompt,
    style,
    use_groq,
    steps,
    seed,
    input_image,
    strength,
    up
):
    if input_image is None:
        return None, None, None, "Input image belum diisi.", history

    pipe = load_pipe("img2img")

    final_prompt = build_prompt(prompt, style, use_groq)

    print("FINAL IMG2IMG PROMPT:", final_prompt)

    seed = normalize_seed(seed)
    input_image = prepare_img2img_input(input_image)

    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=final_prompt,
        negative_prompt=NEGATIVE_PROMPT,
        image=input_image,
        strength=float(strength),
        guidance_scale=7.5,
        num_inference_steps=int(steps),
        generator=generator
    ).images[0]

    if up:
        image = upscale(image)

    history.append(image)

    file_path = save_image(image)

    return image, file_path, seed, final_prompt, history


# =========================
# RETRY HANDLERS
# =========================
def run_txt2img(
    prompt,
    style,
    use_groq,
    steps,
    width,
    height,
    seed,
    up
):
    for _ in range(2):
        try:
            return generate_text_to_image(
                prompt,
                style,
                use_groq,
                steps,
                width,
                height,
                seed,
                up
            )
        except Exception as e:
            print("Retry Text to Image:", e)
            gc.collect()
            time.sleep(2)

    return None, None, None, "Generation failed. Check Space logs.", history


def run_img2img(
    prompt,
    style,
    use_groq,
    steps,
    seed,
    input_image,
    strength,
    up
):
    for _ in range(2):
        try:
            return generate_image_to_image(
                prompt,
                style,
                use_groq,
                steps,
                seed,
                input_image,
                strength,
                up
            )
        except Exception as e:
            print("Retry Image to Image:", e)
            gc.collect()
            time.sleep(2)

    return None, None, None, "Generation failed. Check Space logs.", history


# =========================
# HISTORY FUNCTIONS
# =========================
def select_history(evt: gr.SelectData):
    if evt.index is None:
        return None, None

    if evt.index >= len(history):
        return None, None

    selected_image = history[evt.index]

    return selected_image, evt.index


def reuse_history(selected_index):
    if selected_index is None:
        return None, gr.update(selected="history")

    try:
        selected_index = int(selected_index)
        selected_image = history[selected_index]
    except Exception:
        return None, gr.update(selected="history")

    return selected_image, gr.update(selected="img2img")


def clear_history():
    history.clear()
    gc.collect()

    return [], None, None


def check_groq_status():
    if GROQ_API_KEY:
        return f"✅ Groq API aktif. Model: `{GROQ_MODEL}`"

    return (
        "⚠️ `GROQ_API_KEY` belum ditemukan. "
        "Prompt enhancer akan memakai fallback lokal."
    )


# =========================
# UI
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 AI Image Studio - DreamShaper + Groq Prompt Enhancer")

    gr.Markdown(
        """
        Masukkan prompt dalam **Bahasa Indonesia**.  
        Jika **Groq Prompt Enhancer** aktif, prompt akan diubah otomatis menjadi prompt Bahasa Inggris yang lebih optimal untuk DreamShaper.
        """
    )

    gr.Markdown(check_groq_status())

    selected_history_index = gr.State(value=None)

    with gr.Tabs(selected="txt2img") as tabs:

        # =========================
        # TEXT TO IMAGE TAB
        # =========================
        with gr.Tab("Text to Image", id="txt2img"):
            txt_prompt = gr.Textbox(
                lines=3,
                label="Prompt Bahasa Indonesia",
                placeholder="Contoh: pria memakai jaket hitam berdiri di jalan kota saat hujan malam hari"
            )

            txt_style = gr.Dropdown(
                choices=list(STYLE_PRESETS.keys()),
                value="Realistic",
                label="Style"
            )

            txt_use_groq = gr.Checkbox(
                label="Gunakan Groq Prompt Enhancer",
                value=True
            )

            txt_steps = gr.Slider(
                minimum=10,
                maximum=30,
                value=20,
                step=1,
                label="Steps"
            )

            txt_seed = gr.Number(
                value=-1,
                label="Seed (-1 = random)"
            )

            txt_width = gr.Slider(
                minimum=256,
                maximum=768,
                value=512,
                step=64,
                label="Width"
            )

            txt_height = gr.Slider(
                minimum=256,
                maximum=768,
                value=512,
                step=64,
                label="Height"
            )

            txt_upscale = gr.Checkbox(
                label="Upscale 2x",
                value=False
            )

            txt_button = gr.Button("Generate Text to Image")

            txt_result = gr.Image(
                label="Result"
            )

            txt_seed_output = gr.Number(
                label="Seed Used"
            )

            txt_final_prompt = gr.Textbox(
                lines=5,
                label="Final Prompt ke DreamShaper",
                interactive=False
            )

            txt_file_output = gr.File(
                label="Download Image"
            )

        # =========================
        # IMAGE TO IMAGE TAB
        # =========================
        with gr.Tab("Image to Image", id="img2img"):
            img_prompt = gr.Textbox(
                lines=3,
                label="Prompt Bahasa Indonesia",
                placeholder="Contoh: ubah menjadi cinematic style, lighting dramatis, detail lebih realistis"
            )

            img_input = gr.Image(
                type="pil",
                label="Input Image"
            )

            img_style = gr.Dropdown(
                choices=list(STYLE_PRESETS.keys()),
                value="Realistic",
                label="Style"
            )

            img_use_groq = gr.Checkbox(
                label="Gunakan Groq Prompt Enhancer",
                value=True
            )

            img_strength = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.5,
                step=0.05,
                label="Strength"
            )

            img_steps = gr.Slider(
                minimum=10,
                maximum=30,
                value=20,
                step=1,
                label="Steps"
            )

            img_seed = gr.Number(
                value=-1,
                label="Seed (-1 = random)"
            )

            img_upscale = gr.Checkbox(
                label="Upscale 2x",
                value=False
            )

            img_button = gr.Button("Generate Image to Image")

            img_result = gr.Image(
                label="Result"
            )

            img_seed_output = gr.Number(
                label="Seed Used"
            )

            img_final_prompt = gr.Textbox(
                lines=5,
                label="Final Prompt ke DreamShaper",
                interactive=False
            )

            img_file_output = gr.File(
                label="Download Image"
            )

        # =========================
        # HISTORY TAB
        # =========================
        with gr.Tab("History", id="history"):
            gr.Markdown(
                """
                Gambar yang dibuat dari **Text to Image** atau **Image to Image**
                otomatis muncul di sini.

                Cara memakai ulang:
                1. Klik salah satu gambar di History
                2. Klik tombol **Reuse to Image to Image**
                3. Gambar akan masuk ke menu Image to Image
                """
            )

            history_gallery = gr.Gallery(
                label="History",
                columns=3,
                height=420
            )

            selected_history_image = gr.Image(
                label="Selected Image"
            )

            reuse_button = gr.Button(
                "♻️ Reuse to Image to Image"
            )

            clear_history_button = gr.Button(
                "🗑️ Clear History"
            )

    # =========================
    # EVENTS
    # =========================

    txt_button.click(
        fn=run_txt2img,
        inputs=[
            txt_prompt,
            txt_style,
            txt_use_groq,
            txt_steps,
            txt_width,
            txt_height,
            txt_seed,
            txt_upscale
        ],
        outputs=[
            txt_result,
            txt_file_output,
            txt_seed_output,
            txt_final_prompt,
            history_gallery
        ]
    )

    img_button.click(
        fn=run_img2img,
        inputs=[
            img_prompt,
            img_style,
            img_use_groq,
            img_steps,
            img_seed,
            img_input,
            img_strength,
            img_upscale
        ],
        outputs=[
            img_result,
            img_file_output,
            img_seed_output,
            img_final_prompt,
            history_gallery
        ]
    )

    history_gallery.select(
        fn=select_history,
        outputs=[
            selected_history_image,
            selected_history_index
        ]
    )

    reuse_button.click(
        fn=reuse_history,
        inputs=[
            selected_history_index
        ],
        outputs=[
            img_input,
            tabs
        ]
    )

    clear_history_button.click(
        fn=clear_history,
        inputs=[],
        outputs=[
            history_gallery,
            selected_history_image,
            selected_history_index
        ]
    )

    demo.load(
        fn=lambda: history,
        inputs=[],
        outputs=[
            history_gallery
        ]
    )

# =========================
# ✅ QUEUE SYSTEM
# =========================
demo.queue(max_size=10)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)