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)