Update app.py

app.py

@@ -1,337 +1,516 @@
-import os
-import random
+import gradio as gr
 import numpy as np
+import random
 import torch
-import gradio as gr
-from PIL import Image
-
 import spaces
+from PIL import Image
+from diffusers import QwenImageEditPipeline
+from diffusers.utils import is_xformers_available
+import os
+import base64
+import json
+from huggingface_hub import InferenceClient
+import logging
+
+#############################
+os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
+os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+#############################
+
+def get_caption_language(prompt):
+    """Detects if the prompt contains Chinese characters."""
+    ranges = [
+        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
+    ]
+    for char in prompt:
+        if any(start <= char <= end for start, end in ranges):
+            return 'zh'
+    return 'en'
+
+def polish_prompt(original_prompt, system_prompt, hf_token):
+    """
+    Rewrites the prompt using a Hugging Face InferenceClient.
+    Requires user-provided HF token for API access.
+    """
+    if not hf_token or not hf_token.strip():
+        gr.Warning("HF Token is required for prompt rewriting but was not provided!")
+        return original_prompt
+    client = InferenceClient(
+        provider="cerebras",
+        api_key=hf_token,
+    )
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": original_prompt}
+    ]
+    try:
+        completion = client.chat.completions.create(
+            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
+            messages=messages,
+            max_tokens=512,
+        )
+        polished_prompt = completion.choices[0].message.content
+        polished_prompt = polished_prompt.strip().replace("\n", " ")
+        return polished_prompt
+    except Exception as e:
+        print(f"Error during Hugging Face API call: {e}")
+        gr.Warning("Failed to rewrite prompt. Using original.")
+        return original_prompt
+
+SYSTEM_PROMPT_EDIT = '''
+# Edit Instruction Rewriter
+You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable instruction based on the user's intent and the input image.
+## 1. General Principles
+- Keep the rewritten instruction **concise** and clear.
+- Avoid contradictions, vagueness, or unachievable instructions.
+- Maintain the core logic of the original instruction; only enhance clarity and feasibility.
+- Ensure new added elements or modifications align with the image's original context and art style.
+## 2. Task Types
+### Add, Delete, Replace:
+- When the input is detailed, only refine grammar and clarity.
+- For vague instructions, infer minimal but sufficient details.
+- For replacement, use the format: `"Replace X with Y"`.
+### Text Editing (e.g., text replacement):
+- Enclose text content in quotes, e.g., `Replace "abc" with "xyz"`.
+- Preserving the original structure and language—**do not translate** or alter style.
+### Human Editing (e.g., change a person's face/hair):
+- Preserve core visual identity (gender, ethnic features).
+- Describe expressions in subtle and natural terms.
+- Maintain key clothing or styling details unless explicitly replaced.
+### Style Transformation:
+- If a style is specified, e.g., `Disco style`, rewrite it to encapsulate the essential visual traits.
+- Use a fixed template for **coloring/restoration**:  
+  `"Restore old photograph, remove scratches, reduce noise, enhance details, high resolution, realistic, natural skin tones, clear facial features, no distortion, vintage photo restoration"`  
+  if applicable.
+## 4. Output Format
+Please provide the rewritten instruction in a clean `json` format as:
+{
+  "Rewritten": "..."
+}
+'''
 
-from huggingface_hub import snapshot_download, login
-
-from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_ipadapter import StableDiffusionXLPipeline
-from kolors.models.modeling_chatglm import ChatGLMModel
-from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from kolors.models.unet_2d_condition import UNet2DConditionModel
-from diffusers import AutoencoderKL, EulerDiscreteScheduler
-
-
-# ============= Runtime & Auth =============
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN:
-    login(token=HF_TOKEN)
-    print("Successfully logged in to Hugging Face Hub")
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
-print(f"Device: {DEVICE}, DType: {DTYPE}")
-
-# ============= Weights =============
-# 원본 코드 구조를 따르되, snapshot_download 경로를 그대로 활용합니다.
-print("Downloading models...")
-ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors", token=HF_TOKEN)
-ckpt_dir_ip = snapshot_download(repo_id="Kwai-Kolors/Kolors-IP-Adapter-Plus", token=HF_TOKEN)
-
-# ============= Load Models (IP-Adapter, not FaceID) =============
-# CPU에서는 fp16이 NaN을 유발할 수 있으므로 DTYPE로 통일
-text_encoder = ChatGLMModel.from_pretrained(
-    f"{ckpt_dir}/text_encoder",
-    torch_dtype=DTYPE,
-    trust_remote_code=True
-)
-tokenizer = ChatGLMTokenizer.from_pretrained(
-    f"{ckpt_dir}/text_encoder",
-    trust_remote_code=True
-)
-vae = AutoencoderKL.from_pretrained(
-    f"{ckpt_dir}/vae",
-    revision=None,
-    torch_dtype=DTYPE
-)
-scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
-unet = UNet2DConditionModel.from_pretrained(
-    f"{ckpt_dir}/unet",
-    revision=None,
-    torch_dtype=DTYPE
-)
-
-# CLIP image encoder for IP-Adapter-Plus
-image_encoder = CLIPVisionModelWithProjection.from_pretrained(
-    f"{ckpt_dir_ip}/image_encoder",
-    ignore_mismatched_sizes=True
-).to(dtype=DTYPE, device=DEVICE)
-
-ip_img_size = 336
-clip_image_processor = CLIPImageProcessor(size=ip_img_size, crop_size=ip_img_size)
-
-# StableDiffusionXL pipeline with IP-Adapter (image reference)
-pipe = StableDiffusionXLPipeline(
-    vae=vae,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    unet=unet,
-    scheduler=scheduler,
-    image_encoder=image_encoder,
-    feature_extractor=clip_image_processor,
-    force_zeros_for_empty_prompt=False
-)
-
-# Move core modules to device/dtype
-pipe.vae = pipe.vae.to(DEVICE, dtype=DTYPE)
-pipe.text_encoder = pipe.text_encoder.to(DEVICE, dtype=DTYPE)
-pipe.unet = pipe.unet.to(DEVICE, dtype=DTYPE)
-
-# kolors unet 호환 처리
-if hasattr(pipe.unet, "encoder_hid_proj"):
-    pipe.unet.text_encoder_hid_proj = pipe.unet.encoder_hid_proj
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
 
-# Load IP-Adapter weights (general)
-pipe.load_ip_adapter(
-    f"{ckpt_dir_ip}",
-    subfolder="",
-    weight_name=["ip_adapter_plus_general.bin"]
+# Load LoRA weights for acceleration
+pipe.load_lora_weights(
+    "lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V1.1.safetensors"
 )
+pipe.fuse_lora()
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-def _to_multiple_of_8(x: int) -> int:
-    return int(x // 8 * 8)
-
-
-def _ensure_even(x: int) -> int:
-    return x if x % 2 == 0 else x - 1
-
-
-def _prepare_dims(width: int, height: int) -> tuple[int, int]:
-    # SDXL 권장: 8의 배수 해상도
-    w = _to_multiple_of_8(width)
-    h = _to_multiple_of_8(height)
-    # H.264 등과의 호환성을 고려해 짝수 유지(선택)
-    w = _ensure_even(w)
-    h = _ensure_even(h)
-    return max(256, min(MAX_IMAGE_SIZE, w)), max(256, min(MAX_IMAGE_SIZE, h))
-
-
-def _move_to_device():
-    # 호출 시점에 안전하게 보장
-    global pipe, image_encoder
-    pipe.vae = pipe.vae.to(DEVICE, dtype=DTYPE)
-    pipe.text_encoder = pipe.text_encoder.to(DEVICE, dtype=DTYPE)
-    pipe.unet = pipe.unet.to(DEVICE, dtype=DTYPE)
-    image_encoder = image_encoder.to(device=DEVICE, dtype=DTYPE)
-    pipe.image_encoder = image_encoder
-
-
-def _generate(
-    prompt: str,
-    ip_adapter_image: Image.Image,
-    ip_adapter_scale: float,
-    negative_prompt: str,
-    seed: int,
-    width: int,
-    height: int,
-    guidance_scale: float,
-    num_inference_steps: int,
-):
-    _move_to_device()
-    pipe.set_ip_adapter_scale([ip_adapter_scale])
-
-    # 해상도 정규화
-    width, height = _prepare_dims(width, height)
-
-    generator = torch.Generator(device=DEVICE).manual_seed(seed)
-
-    with torch.no_grad():
-        if DEVICE == "cuda":
-            with torch.autocast(device_type="cuda", dtype=torch.float16):
-                images = pipe(
-                    prompt=prompt,
-                    ip_adapter_image=[ip_adapter_image],
-                    negative_prompt=negative_prompt,
-                    height=height,
-                    width=width,
-                    num_inference_steps=int(num_inference_steps),
-                    guidance_scale=float(guidance_scale),
-                    num_images_per_prompt=1,
-                    generator=generator,
-                ).images
-        else:
-            images = pipe(
-                prompt=prompt,
-                ip_adapter_image=[ip_adapter_image],
-                negative_prompt=negative_prompt,
-                height=height,
-                width=width,
-                num_inference_steps=int(num_inference_steps),
-                guidance_scale=float(guidance_scale),
-                num_images_per_prompt=1,
-                generator=generator,
-            ).images
-
-    return images[0]
-
+if is_xformers_available():
+    pipe.enable_xformers_memory_efficient_attention()
+else:
+    print("xformers not available or failed to load.")
 
-# Spaces GPU 스케줄러는 CUDA가 있을 때만 감쌉니다.
-def _infer_core(
+@spaces.GPU(duration=60)
+def infer(
+    image,
     prompt,
-    ip_adapter_image,
-    ip_adapter_scale=0.5,
-    negative_prompt="",
-    seed=100,
+    seed=42,
     randomize_seed=False,
-    width=1024,
-    height=1024,
-    guidance_scale=5.0,
-    num_inference_steps=50,
+    true_guidance_scale=1.0,
+    num_inference_steps=8,
+    rewrite_prompt=False,
+    hf_token="",
+    num_images_per_prompt=1,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if ip_adapter_image is None:
-        gr.Warning("Please upload an IP-Adapter reference image.")
-        return None, 0
-
+    """
+    Requires user-provided HF token for prompt rewriting.
+    """
+    original_prompt = prompt  # Save original prompt for display
+    negative_prompt = " "
+    prompt_info = ""  # Initialize info text
+    
+    # Handle prompt rewriting with status messages
+    if rewrite_prompt:
+        if not hf_token.strip():
+            gr.Warning("HF Token is required for prompt rewriting but was not provided!")
+            prompt_info = f"""<div class="prompt-info-box warning">
+<h3>⚠️ Prompt Rewriting Skipped</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p class="note">HF Token required for enhancement</p>
+</div>"""
+            rewritten_prompt = original_prompt
+        else:
+            try:
+                rewritten_prompt = polish_prompt(original_prompt, SYSTEM_PROMPT_EDIT, hf_token)
+                prompt_info = f"""<div class="prompt-info-box success">
+<h3>✨ Enhanced Successfully</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p><strong>Enhanced:</strong> {rewritten_prompt}</p>
+</div>"""
+            except Exception as e:
+                gr.Warning(f"Prompt rewriting failed: {str(e)}")
+                rewritten_prompt = original_prompt
+                prompt_info = f"""<div class="prompt-info-box error">
+<h3>❌ Enhancement Failed</h3>
+<p><strong>Original:</strong> {original_prompt}</p>
+<p class="note">Error: {str(e)}</p>
+</div>"""
+    else:
+        rewritten_prompt = original_prompt
+        prompt_info = f"""<div class="prompt-info-box default">
+<h3>📝 Original Prompt</h3>
+<p>{original_prompt}</p>
+</div>"""
+    
+    # Generate images
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    edited_images = pipe(
+        image,
+        prompt=rewritten_prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        true_cfg_scale=true_guidance_scale,
+        num_images_per_prompt=num_images_per_prompt,
+    ).images
+    
+    return edited_images, seed, prompt_info
 
-    image = _generate(
-        prompt=prompt or "",
-        ip_adapter_image=ip_adapter_image,
-        ip_adapter_scale=float(ip_adapter_scale),
-        negative_prompt=negative_prompt or "",
-        seed=int(seed),
-        width=int(width),
-        height=int(height),
-        guidance_scale=float(guidance_scale),
-        num_inference_steps=int(num_inference_steps),
-    )
-    return image, seed
-
-
-if torch.cuda.is_available():
-    infer = spaces.GPU(duration=80)(_infer_core)
-else:
-    infer = _infer_core
-
+MAX_SEED = np.iinfo(np.int32).max
 
 examples = [
-    ["A dog", "minta.jpeg", 0.4],
-    ["A capybara", "king-min.png", 0.5],
-    ["A cat", "blue_hair.png", 0.5],
-    ["", "meow.jpeg", 1.0],
+    "Replace the cat with a friendly golden retriever. Make it look happier, and add more background details.",
+    "Add text 'Qwen - AI for image editing' in Chinese at the bottom center with a small shadow.",
+    "Change the style to 1970s vintage, add old photo effect, restore any scratches on the wall or window.",
+    "Remove the blue sky and replace it with a dark night cityscape.",
+    """Replace "Qwen" with "通义" in the Image. Ensure Chinese font is used for "通义" and position it to the top left with a light heading-style font."""
 ]
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 720px;
+# Custom CSS for enhanced visual design
+custom_css = """
+/* Gradient background */
+.gradio-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 25%, #f093fb 50%, #fecfef 75%, #fecfef 100%);
+    min-height: 100vh;
+}
+/* Main container styling */
+.container {
+    max-width: 1400px !important;
+    margin: 0 auto !important;
+    padding: 2rem !important;
+}
+/* Card-like sections */
+.gr-box {
+    background: rgba(255, 255, 255, 0.95) !important;
+    backdrop-filter: blur(10px) !important;
+    border-radius: 20px !important;
+    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1) !important;
+    border: 1px solid rgba(255, 255, 255, 0.5) !important;
+    padding: 1.5rem !important;
+    margin-bottom: 1.5rem !important;
+}
+/* Header styling */
+h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    font-size: 3rem !important;
+    font-weight: 800 !important;
+    text-align: center;
+    margin-bottom: 0.5rem !important;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+}
+h2 {
+    color: #4a5568 !important;
+    font-size: 1.5rem !important;
+    font-weight: 600 !important;
+    margin-bottom: 1rem !important;
+}
+/* Button styling */
+.gr-button-primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    color: white !important;
+    font-weight: 600 !important;
+    font-size: 1.1rem !important;
+    padding: 0.8rem 2rem !important;
+    border-radius: 12px !important;
+    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
+    transition: all 0.3s ease !important;
+}
+.gr-button-primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.5) !important;
+}
+/* Input fields styling */
+.gr-input, .gr-text-input, .gr-slider, .gr-dropdown {
+    border-radius: 10px !important;
+    border: 2px solid #e2e8f0 !important;
+    background: white !important;
+    transition: all 0.3s ease !important;
+}
+.gr-input:focus, .gr-text-input:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
+/* Accordion styling */
+.gr-accordion {
+    background: rgba(255, 255, 255, 0.8) !important;
+    border-radius: 12px !important;
+    border: 1px solid rgba(102, 126, 234, 0.2) !important;
+    overflow: hidden !important;
+}
+/* Gallery styling */
+.gr-gallery {
+    border-radius: 12px !important;
+    overflow: hidden !important;
+}
+/* Prompt info boxes */
+.prompt-info-box {
+    padding: 1.5rem;
+    border-radius: 12px;
+    margin: 1rem 0;
+    animation: fadeIn 0.5s ease;
+}
+.prompt-info-box h3 {
+    margin: 0 0 0.75rem 0;
+    font-size: 1.2rem;
+    font-weight: 600;
+}
+.prompt-info-box p {
+    margin: 0.5rem 0;
+    line-height: 1.6;
+}
+.prompt-info-box.success {
+    background: linear-gradient(135deg, #d4f4dd 0%, #e3f9e5 100%);
+    border-left: 4px solid #48bb78;
+}
+.prompt-info-box.warning {
+    background: linear-gradient(135deg, #fef5e7 0%, #fff9ec 100%);
+    border-left: 4px solid #f6ad55;
+}
+.prompt-info-box.error {
+    background: linear-gradient(135deg, #fed7d7 0%, #fee5e5 100%);
+    border-left: 4px solid #fc8181;
 }
-#result img{
-    object-position: top;
+.prompt-info-box.default {
+    background: linear-gradient(135deg, #e6f3ff 0%, #f0f7ff 100%);
+    border-left: 4px solid #667eea;
 }
-#result .image-container{
-    height: 100%
+.prompt-info-box .note {
+    font-size: 0.9rem;
+    color: #718096;
+    font-style: italic;
+}
+/* Checkbox styling */
+.gr-checkbox {
+    background: white !important;
+    border-radius: 8px !important;
+    padding: 0.5rem !important;
+}
+/* Token input field */
+input[type="password"] {
+    font-family: monospace !important;
+    letter-spacing: 0.05em !important;
+}
+/* Info badges */
+.gr-markdown p {
+    color: #4a5568;
+    line-height: 1.6;
+}
+.gr-markdown a {
+    color: #667eea !important;
+    text-decoration: none !important;
+    font-weight: 500 !important;
+    transition: color 0.3s ease !important;
+}
+.gr-markdown a:hover {
+    color: #764ba2 !important;
+    text-decoration: underline !important;
+}
+/* Animation */
+@keyframes fadeIn {
+    from {
+        opacity: 0;
+        transform: translateY(10px);
+    }
+    to {
+        opacity: 1;
+        transform: translateY(0);
+    }
+}
+/* Slider styling */
+.gr-slider input[type="range"] {
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%) !important;
+}
+/* Group styling */
+.gr-group {
+    background: rgba(249, 250, 251, 0.8) !important;
+    border-radius: 12px !important;
+    padding: 1rem !important;
+    margin-top: 1rem !important;
+}
+/* Loading spinner customization */
+.gr-loading {
+    color: #667eea !important;
+}
+/* Example buttons */
+.gr-examples button {
+    background: white !important;
+    border: 2px solid #e2e8f0 !important;
+    border-radius: 8px !important;
+    padding: 0.5rem 1rem !important;
+    transition: all 0.3s ease !important;
+}
+.gr-examples button:hover {
+    border-color: #667eea !important;
+    background: rgba(102, 126, 234, 0.05) !important;
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown("# Kolors IP-Adapter - image reference and variations")
 
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Run", scale=0)
-
-        with gr.Row():
-            with gr.Column():
-                ip_adapter_image = gr.Image(label="IP-Adapter Image", type="pil")
-                ip_adapter_scale = gr.Slider(
-                    label="Image influence scale",
-                    info="Use 1 for creating variations",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.05,
-                    value=0.5,
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎨 Nano-Banana")
+    gr.Markdown("✨ **Ultra-fast 8-step image editing with AI-powered prompt enhancement**")
+    gr.Markdown("🔐 **Secure prompt rewriting with your [Hugging Face token](https://huggingface.co/settings/tokens)**")
+    
+    # 배지를 가운데 정렬하여 나란히 배치
+    gr.HTML("""
+        <div style="display: flex; justify-content: center; align-items: center; gap: 20px; margin: 20px 0;">
+            <a href="https://huggingface.co/spaces/Heartsync/Nano-Banana" target="_blank">
+                <img src="https://img.shields.io/static/v1?label=OPEN%20NANO-BANANA&message=Image%20EDITOR&color=%230000ff&labelColor=%23800080&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
+            </a>
+        </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group():
+                input_image = gr.Image(
+                    label="📸 Input Image", 
+                    type="pil",
+                    elem_classes="gr-box"
                 )
-            result = gr.Image(label="Result", elem_id="result")
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-            )
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                prompt = gr.Text(
+                    label="✏️ Edit Instruction", 
+                    placeholder="e.g. Add a dog to the right side, change the sky to sunset...",
+                    lines=3,
+                    elem_classes="gr-box"
                 )
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
+                
+                with gr.Accordion("⚙️ Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0
+                    )
+                    randomize_seed = gr.Checkbox(label="🎲 Randomize Seed", value=True)
+                    
+                    with gr.Row():
+                        true_guidance_scale = gr.Slider(
+                            label="Guidance Scale",
+                            minimum=1.0,
+                            maximum=5.0,
+                            step=0.1,
+                            value=4.0
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Inference Steps",
+                            minimum=4,
+                            maximum=16,
+                            step=1,
+                            value=8
+                        )
+                    
+                    num_images_per_prompt = gr.Slider(
+                        label="Images per Prompt",
+                        minimum=1,
+                        maximum=4,
+                        step=1,
+                        value=1
+                    )
+                
+                run_button = gr.Button("🚀 Generate Edit", variant="primary", size="lg")
+        
+        with gr.Column(scale=1):
+            result = gr.Gallery(
+                label="🖼️ Output Images", 
+                show_label=True, 
+                columns=2,
+                rows=2,
+                elem_classes="gr-box"
+            )
+            
+            # Prompt display component
+            prompt_info = gr.HTML(visible=False)
+            
+            with gr.Group():
+                rewrite_toggle = gr.Checkbox(
+                    label="🤖 Enable AI Prompt Enhancement", 
+                    value=False, 
+                    interactive=True
                 )
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=5.0,
+                hf_token_input = gr.Textbox(
+                    label="🔑 Hugging Face API Token",
+                    type="password",
+                    placeholder="hf_xxxxxxxxxxxxxxxx",
+                    visible=False,
+                    info="Your token is secure and only used for API calls. Get yours from HuggingFace settings.",
+                    elem_classes="gr-box"
                 )
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=25,
+                
+                def toggle_token_visibility(checked):
+                    return gr.update(visible=checked)
+                
+                rewrite_toggle.change(
+                    toggle_token_visibility,
+                    inputs=[rewrite_toggle],
+                    outputs=[hf_token_input]
                 )
-
-        # 파일 예시가 로컬에 없을 수 있어 cache_examples="lazy" 유지
-        gr.Examples(
-            examples=examples,
-            fn=infer,
-            inputs=[prompt, ip_adapter_image, ip_adapter_scale],
-            outputs=[result, seed],
-            cache_examples="lazy",
-        )
-
+    
+    # Examples section
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        label="💡 Example Prompts"
+    )
+    
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            input_image,
             prompt,
-            ip_adapter_image,
-            ip_adapter_scale,
-            negative_prompt,
             seed,
             randomize_seed,
-            width,
-            height,
-            guidance_scale,
+            true_guidance_scale,
             num_inference_steps,
+            rewrite_toggle,
+            hf_token_input,
+            num_images_per_prompt
         ],
-        outputs=[result, seed],
+        outputs=[result, seed, prompt_info]
+    )
+    
+    # Show prompt info box after processing
+    def set_prompt_visible():
+        return gr.update(visible=True)
+    
+    run_button.click(
+        fn=set_prompt_visible,
+        inputs=None,
+        outputs=[prompt_info],
+        queue=False
+    )
+    prompt.submit(
+        fn=set_prompt_visible,
+        inputs=None,
+        outputs=[prompt_info],
+        queue=False
     )
 
-demo.queue().launch()
+if __name__ == "__main__":
+    demo.launch()