ui.py

"""
Gradio UI für AI Image Generator - Mit Mock-Funktionen für CPU-Testing
"""
import gradio as gr
from PIL import Image, ImageDraw
import random
import time

# ==================== KONSTANTEN ====================
MAX_IMAGE_SIZE = 4096  # Maximale Bildgröße für Verarbeitung

# ==================== MOCK-FUNKTIONEN ====================
def mock_text_to_image(prompt, model_id, steps, guidance_scale, progress=None):
    """Mock für Text-zu-Bild: Erzeugt einfaches Farbbild"""
    print(f"📝 Mock text_to_image aufgerufen: '{prompt[:50]}...'")
    
    if progress:
        for i in range(10):
            time.sleep(0.05)
            progress((i+1)/10, desc="Mock-Generierung läuft...")
    
    colors = ["lightblue", "lightgreen", "lavender", "peachpuff"]
    img = Image.new('RGB', (512, 512), color=random.choice(colors))
    
    draw = ImageDraw.Draw(img)
    text = f"Mock: {prompt[:30]}..." if len(prompt) > 30 else f"Mock: {prompt}"
    draw.text((50, 256), text, fill="black")
    
    status = f"✅ Mock: '{prompt[:30]}...' würde mit {model_id} generiert (Steps: {steps}, CFG: {guidance_scale})"
    return img, status

def mock_img_to_image(image, prompt, neg_prompt, strength, steps, guidance_scale,
                     mode, bbox_x1, bbox_y1, bbox_x2, bbox_y2, progress=None):
    """Mock für Bild-zu-Bild: Fügt Rahmen und Text hinzu"""
    print(f"🎨 Mock img_to_image aufgerufen: Modus={mode}, Prompt='{prompt[:30]}...'")
    
    if not image:
        return None, None, None, None, None
    
    if progress:
        for i in range(10):
            time.sleep(0.05)
            progress((i+1)/10, desc=f"Mock-{mode} läuft...")
    
    img = image.copy().convert("RGB")
    draw = ImageDraw.Draw(img)
    
    colors = {
        "environment_change": "green",
        "focus_change": "orange", 
        "face_only_change": "red"
    }
    border_color = colors.get(mode, "blue")
    
    draw.rectangle([20, 20, img.width-20, img.height-20], 
                   outline=border_color, width=10)
    
    if all(v is not None for v in [bbox_x1, bbox_y1, bbox_x2, bbox_y2]):
        x1, y1 = min(bbox_x1, bbox_x2), min(bbox_y1, bbox_y2)
        x2, y2 = max(bbox_x1, bbox_x2), max(bbox_y1, bbox_y2)
        draw.rectangle([x1, y1, x2, y2], outline="yellow", width=5)
        draw.text((x1+5, y1+5), f"{mode}", fill="white")
    
    draw.text((30, 30), f"Mock: {mode}", fill=border_color)
    draw.text((30, 60), f"Strength: {strength}", fill="black")
    
    mask_preview = Image.new('RGB', (256, 256), color='gray')
    controlnet_map = Image.new('RGB', (256, 256), color='darkblue')
    canny_map = Image.new('RGB', (256, 256), color='black')
    
    status = f"✅ Mock: {mode} angewendet (Strength: {strength}, Steps: {steps})"
    print(status)
    
    return img, mask_preview, mask_preview, controlnet_map, canny_map

def sort_coordinates(x1, y1, x2, y2):
    """Sortiert Koordinaten, so dass x1 <= x2 und y1 <= y2"""
    sorted_x1 = min(x1, x2)
    sorted_x2 = max(x1, x2)
    sorted_y1 = min(y1, y2)
    sorted_y2 = max(y1, y2)
    return sorted_x1, sorted_y1, sorted_x2, sorted_y2

def create_preview_image(image, bbox_coords, mode):
    """Vorschau mit dynamischer Rahmendicke"""
    if image is None:
        return None
        
    preview = image.copy()
    draw = ImageDraw.Draw(preview)
    
    if mode == "environment_change":
        border_color = (0, 255, 0, 180)
        mode_text = "UMGEBUNG ÄNDERN"
        box_color = (255, 255, 0, 200)
        text_bg_color = (0, 128, 0, 160)
    elif mode == "focus_change":
        border_color = (255, 165, 0, 180)
        mode_text = "FOCUS VERÄNDERN"
        box_color = (255, 0, 0, 200)
        text_bg_color = (255, 140, 0, 160)
    elif mode == "face_only_change":
        border_color = (255, 0, 0, 180)
        mode_text = "NUR BEREICH"
        box_color = (255, 0, 0, 200)
        text_bg_color = (128, 0, 0, 160)
    else:
        border_color = (128, 128, 128, 180)
        mode_text = "UNBEKANNT"
        box_color = (128, 128, 128, 200)
        text_bg_color = (64, 64, 64, 160)
    
    border_width = max(8, image.width // 200)
    box_width = max(3, image.width // 400)
    
    draw.rectangle([0, 0, preview.width-1, preview.height-1], 
                  outline=border_color, width=border_width)
    
    if bbox_coords and all(coord is not None for coord in bbox_coords):
        x1, y1, x2, y2 = sort_coordinates(*bbox_coords)
        x1 = max(0, min(x1, preview.width-1))
        y1 = max(0, min(y1, preview.height-1))
        x2 = max(0, min(x2, preview.width-1))
        y2 = max(0, min(y2, preview.height-1))
        
        if x2 > x1 and y2 > y1:
            draw.rectangle([x1, y1, x2, y2], outline=box_color, width=box_width)
            text_color = (255, 255, 255)
            text_y = max(0, y1 - 25)
            
            try:
                from PIL import ImageFont
                font_size = max(12, image.width // 50)
                font = ImageFont.truetype("arial.ttf", font_size)
                text_bbox = draw.textbbox((x1, text_y), mode_text, font=font)
                draw.rectangle([text_bbox[0]-5, text_bbox[1]-2, 
                              text_bbox[2]+5, text_bbox[3]+2], 
                             fill=text_bg_color)
                draw.text((x1, text_y), mode_text, fill=text_color, font=font)
            except:
                text_bbox = draw.textbbox((x1, text_y), mode_text)
                draw.rectangle([text_bbox[0]-5, text_bbox[1]-2, 
                              text_bbox[2]+5, text_bbox[3]+2], 
                             fill=text_bg_color)
                draw.text((x1, text_y), mode_text, fill=text_color)
    
    return preview

def mock_update_live_preview(image, bbox_x1, bbox_y1, bbox_x2, bbox_y2, mode):
    """Mock für Live-Vorschau MIT dynamischen Rahmen"""
    if not image:
        return None
    
    bbox_coords = sort_coordinates(bbox_x1, bbox_y1, bbox_x2, bbox_y2)
    return create_preview_image(image, bbox_coords, mode)

def mock_process_image_upload(image):
    """Mock für Bild-Upload: Setzt BBox in der Mitte"""
    if not image:
        return None, 100, 100, 300, 300  # x1, y1, x2, y2
    
    w, h = image.size
    bbox_size = min(w, h) * 0.3
    x1 = (w - bbox_size) / 2
    y1 = (h - bbox_size) / 4
    x2 = x1 + bbox_size
    y2 = y1 + bbox_size * 1.2
    
    print(f"📐 Bildgröße erkannt: {w}x{h} -> BBox: {int(x1)}-{int(x2)}, {int(y1)}-{int(y2)}")
    
    preview = mock_update_live_preview(image, x1, y1, x2, y2, "environment_change")
    return preview, int(x1), int(y1), int(x2), int(y2)  # x1, y1, x2, y2

def mock_update_slider_for_image(image):
    """KORRIGIERT: Slider-Update mit getrennten Maxima für Breite und Höhe"""
    if not image:
        return (
            gr.update(maximum=MAX_IMAGE_SIZE),
            gr.update(maximum=MAX_IMAGE_SIZE),
            gr.update(maximum=MAX_IMAGE_SIZE),
            gr.update(maximum=MAX_IMAGE_SIZE)
        )
    
    w, h = image.size
    max_x = min(w, MAX_IMAGE_SIZE)
    max_y = min(h, MAX_IMAGE_SIZE)
    
    print(f"📐 Slider-Maxima gesetzt: X={max_x}, Y={max_y} (Bild: {w}x{h})")
    
    return (
        gr.update(maximum=max_x),  # bbox_x1: linke Kante max = Bildbreite
        gr.update(maximum=max_y),  # bbox_y1: obere Kante max = Bildhöhe
        gr.update(maximum=max_x),  # bbox_x2: rechte Kante max = Bildbreite
        gr.update(maximum=max_y),  # bbox_y2: untere Kante max = Bildhöhe
    )

def mock_update_model_settings(model_id):
    configs = {
        "runwayml/stable-diffusion-v1-5": (35, 7.5, "🏠 SD 1.5 Mock"),
        "SG161222/Realistic_Vision_V6.0_B1_noVAE": (40, 7.0, "👤 Realistic Vision Mock")
    }
    steps, cfg, msg = configs.get(model_id, (35, 7.5, "Standard Mock"))
    # WICHTIG: Die dritte Rückgabe muss der HTML-String für die Info-Box sein.
    info_html = f"<div class='model-info-box'><strong>{msg}</strong><br>Mock-Einstellungen: {steps} Steps, CFG {cfg}</div>"
    return steps, cfg, info_html


def mock_update_info(mode):
    """Mock für Prompt-Info-Update - ALTE TEXTE WIEDERHERGESTELLT"""
    if mode == "environment_change":
        return (
            "`[STIL-MOTIV],[UMGEBUNG],[PERSPEKTIVE],[DETAILS],[QUALITÄT],[BELEUCHTUNG]`",
            "`[GESICHTER/ANATOMIE], [FEHLER], [QUALITÄT], [UNERWÜNSCHTES]`"
        )
    elif mode == "focus_change":
        return (
            "`[GESICHTSBESCHREIBUNG], [KLEIDUNG], [POSITION], [DETAILS], [STIL]`",
            "`[DEFORMIERT], [UNSCHÄRFE], [ANATOMIEFEHLER], [UNERWÜNSCHTES]`"
        )
    else:
        return (
            "`[STIL/KOPFART],[HAARFARBE],[AUGEN],[GESICHTSAUSDRUCK],[DETAILS],[BELEUCHTUNG]`",
            "`[UNREALISTISCH], [ASYMETRISCH], [FEHLER], [UNERWÜNSCHTES]`"
        )

# ==================== UI-DEFINITION ====================
def main_ui():
    """Haupt-UI-Funktion (angepasst für 3 Modi)"""
    with gr.Blocks(
        title="AI Image Generator - Mock Version",
        theme=gr.themes.Base(),
        css="""
/* ===== INFO-BOXEN MIT GRÖßEREM TEXT ===== */
.info-box {
    background: #f8fafc;
    padding: 4px 4px;
    border-radius: 4px;
    border: none;
    margin-bottom: 3px;
    font-size: 12px;
    line-height: 1.3;
    height: 40px;       <!-- FESTE Höhe statt auto
    min-height: 40px;
    color: #475569;
    display: flex;       <!-- FLEX statt block!
    align-items: center;  <!-- Jetzt funktioniert vertikale Zentrierung
    justify-content: center;    <!-- Horizontale Zentrierung
    text-align: center;
    font-family: 'Segoe UI', 'Monaco', monospace;
    overflow: hidden;      <!-- Bei zu langem Text
    white-space: normal;
    word-wrap: break-word;
}

/* Code in Info-Boxen größer */
.info-box code {
    background: #ffffff;
    padding: 2px 2px;
    border-radius: 3px;
    font-size: 12px;
    border: none;
    color: #334155;
    font-weight: 500;
}

/* ===== DICKE SCHWARZE LABELS ===== */
.gr-column:first-child .gr-box .wrap .title {
    color: #000000;
    font-weight: 900;
    font-size: 16px;
}

.gr-column:last-child .gr-box .wrap .title {
    color: #000000;
    font-weight: 900;
    font-size: 16px;
}

/* ===== TEXTBOXEN (FESTE HÖHE) ===== */
.prompt-box textarea {
    min-height: 90px !important;
    max-height: 90px !important;
    height: 90px !important;
    border-radius: 6px !important;
    border: 2px solid #3b82f6 !important;
    padding: 12px !important;
    font-size: 14px !important;
    background: white !important;
    box-shadow: 0 1px 3px rgba(0,0,0,0.05) !important;
    resize: vertical !important;
    overflow-y: auto !important;
}

.prompt-box textarea:focus {
    border-color: #1d4ed8 !important;
    box-shadow: 0 0 0 3px rgba(59, 130, 246, 0.1) !important;
}

/* ===== VISUELLE HIERARCHIE ===== */
.gr-column:first-child .info-box {
    background: #f0f9ff !important;  /* KEIN border-left mehr! */
}

.gr-column:last-child .info-box {
    background: #fef2f2 !important;  /* KEIN border-left mehr! */
}

.model-info-box { 
    background: #e8f4fd; 
    padding: 12px; 
    border-radius: 6px; 
    margin: 10px 0; 
    border-left: 4px solid #2196f3; 
    font-size: 14px !important;
}

#generate-button { 
    background-color: #0080FF !important; 
    margin: 20px auto !important; 
    width: 280px; 
    font-size: 16px !important;
    font-weight: 600 !important;
}

.radio-group { 
    background: #f8f9fa; 
    padding: 15px; 
    border-radius: 8px; 
    margin: 10px 0; 
    font-size: 14px !important;
}

.status-message {
    padding: 10px;
    border-radius: 5px;
    margin: 10px 0;
    font-size: 14px;
    background: #f0f9ff;
    border-left: 4px solid #3b82f6;
}

/* Slider besser sichtbar */
.gr-slider {
    margin: 8px 0 !important;
}

.gr-slider .wrap {
    padding: 8px 0 !important;
}
"""
    ) as demo:
        
        with gr.Tab("Text zu Bild"):
            gr.Markdown("## 🎨 Text zu Bild Generator (Mock Version)")
            
            with gr.Row():
                with gr.Column(scale=2):
                    model_dropdown = gr.Dropdown(
                        choices=[
                            ("🏠 Stable Diffusion 1.5 (Mock)", "runwayml/stable-diffusion-v1-5"),
                            ("👤 Realistic Vision V6.0 (Mock)", "SG161222/Realistic_Vision_V6.0_B1_noVAE")
                        ],
                        value="runwayml/stable-diffusion-v1-5",
                        label="📁 Modellauswahl"
                    )
                    
                    model_info_box = gr.Markdown(                       
                        value="<div class='model-info-box'><strong>🏠 Stable Diffusion 1.5 (Mock)</strong><br>Universal model, good all-rounder<br><em>Mock-Einstellungen: 35 Steps, CFG 7.5</em></div>",
                        label="Modellinformationen"
                    )
                    
                with gr.Column(scale=3):
                    txt_input = gr.Textbox(
                        placeholder="z.B. ultra realistic mountain landscape at sunrise...",
                        lines=3,
                        label="🎯 Prompt (Englisch)",
                        #value="beautiful landscape sunset mountains",
                        elem_classes=["prompt-box"]
                    )
            
            with gr.Row():
                with gr.Column():
                    txt_steps = gr.Slider(
                        minimum=10, maximum=100, value=35, step=1,
                        label="⚙️ Inferenz-Schritte"
                    )
                with gr.Column():
                    txt_guidance = gr.Slider(
                        minimum=1.0, maximum=20.0, value=7.5, step=0.5,
                        label="🎛️ Prompt-Stärke (CFG Scale)"
                    )

            status_output = gr.Markdown(
                value="<div class='status-message'>ℹ️ Wählen Sie ein Modell und geben Sie einen Prompt ein.</div>",
                elem_classes="status-message"
            )
            
            generate_btn = gr.Button("🚀 Mock-Bild generieren", variant="primary", elem_id="generate-button")
            
            with gr.Row():
                txt_output = gr.Image(
                    label="🖼️ Generiertes Mock-Bild", 
                    show_download_button=True,
                    type="pil",
                    height=400
                )
            
            model_dropdown.change(
                fn=mock_update_model_settings,
                inputs=[model_dropdown],
                outputs=[txt_steps, txt_guidance, model_info_box]
            )
            
            generate_btn.click(
                fn=mock_text_to_image,
                inputs=[txt_input, model_dropdown, txt_steps, txt_guidance],
                outputs=[txt_output, status_output],
                concurrency_limit=1
            )
        
        with gr.Tab("Bild zu Bild"):
            gr.Markdown("## 🖼️ Bild zu Bild Transformation (Mock Version)")
            
            with gr.Row():
                with gr.Column():
                    img_input = gr.Image(
                        type="pil", 
                        label="📤 Eingabebild",
                        height=300,
                        sources=["upload"],
                        elem_id="image-upload"
                    )
                with gr.Column():
                    preview_output = gr.Image(
                        label="Vorschau",
                        height=300,
                        interactive=False
                    )
            
            with gr.Row():
                with gr.Column():
                    mode_radio = gr.Radio(
                        choices=[
                            ("🌳 Umgebung ändern", "environment_change"),
                            ("🎯 Focus verändern", "focus_change"),
                            ("👤 Ausschließlich Gesicht", "face_only_change")
                        ],
                        value="environment_change",
                        label="Wähle den Transformationsmodus:",
                        elem_classes="radio-group"
                    )
            
            with gr.Row():
                gr.Markdown("### 📐 Bildelementbereich anpassen")
            
    
            with gr.Row():
                with gr.Column():
                    bbox_x1 = gr.Slider(
                        label="← Links (x1)", 
                        minimum=0, maximum=MAX_IMAGE_SIZE, value=100, step=1,
                        info="Linke Kante des Bildelementbereichs"
                    )
                with gr.Column():
                    bbox_y1 = gr.Slider(
                        label="↑ Oben (y1)", 
                        minimum=0, maximum=MAX_IMAGE_SIZE, value=100, step=1,
                        info="Obere Kante des Bildelementbereichs"
                    )
            
            with gr.Row():
                with gr.Column():
                    bbox_x2 = gr.Slider(
                        label="→ Rechts (x2)", 
                        minimum=0, maximum=MAX_IMAGE_SIZE, value=300, step=1,
                        info="Rechte Kante des Bildelementbereichs"
                    )
                with gr.Column():
                    bbox_y2 = gr.Slider(
                        label="↓ Unten (y2)", 
                        minimum=0, maximum=MAX_IMAGE_SIZE, value=300, step=1,
                        info="Untere Kante des Bildelementbereichs"
                    )
            
            with gr.Row():
                with gr.Column():
                    pos_info = gr.Markdown(
                        value="`[STIL-MOTIV],[UMGEBUNG],[PERSPEKTIVE],[DETAILS],[QUALITÄT],[BELEUCHTUNG]`",
                        elem_classes=["info-box"]
                    )
                    img_prompt = gr.Textbox(
                        placeholder="photorealistic coastal beach, keep person unchanged...",
                        lines=2,
                        #label="<span style='color: black; font-weight: 900; font-size: 16px;'>🎯 Transformations-Prompt</span>",
                        label="🎯 Transformations-Prompt",
                        elem_classes=["prompt-box"]
                    )
                with gr.Column():
                    neg_info = gr.Markdown(
                        value="`[GESICHTER/ANATOMIE], [FEHLER], [QUALITÄT], [UNERWÜNSCHTES]`",
                        elem_classes=["info-box"]
                    )        
                    img_neg_prompt = gr.Textbox(
                        placeholder="blurry face, deformed anatomy...",
                        lines=2,
                        #label="<span style='color: black; font-weight: 900; font-size: 16px;'>🚫 Negativ-Prompt</span>",
                        label="🚫 Negativ-Prompt",
                        elem_classes=["prompt-box"]
                    )
            
            with gr.Row():
                with gr.Column():
                    strength_slider = gr.Slider(
                        minimum=0.1, maximum=0.9, value=0.4, step=0.05,
                        label="💪 Veränderungs-Stärke (strength)"
                    )
                with gr.Column():
                    img_steps = gr.Slider(
                        minimum=10, maximum=45, value=35, step=1,
                        label="⚙️ Inferenz-Schritte"
                    )
                with gr.Column():
                    img_guidance = gr.Slider(
                        minimum=1.0, maximum=15.0, value=7.5, step=0.5,
                        label="🎛️ Prompt-Stärke (guidance)"
                    )
            
            transform_btn = gr.Button("🔄 Mock-Bild transformieren", variant="primary")
            
            with gr.Row():
                img_output = gr.Image(
                    label="✨ Transformiertes Mock-Bild",
                    show_download_button=True,
                    type="pil",
                    height=400
                )
            
            with gr.Row():
                sam_raw_mask_output = gr.Image(
                    label="🔍 SAM-Rohmaske (Mock)",
                    type="pil",
                    height=250,
                    show_download_button=False
                )
                processed_mask_output = gr.Image(
                    label="🛠️ Nachbearbeitete Maske (Mock)",
                    type="pil",
                    height=250,
                    show_download_button=False
                )
            
            with gr.Row():                 
                pose_map_output = gr.Image(
                    label="🎭 Pose/Depth Map (Mock)", 
                    type="pil", 
                    height=250,
                    show_download_button=False
                )
                canny_map_output = gr.Image(
                    label="📐 Canny Edge Map (Mock)", 
                    type="pil", 
                    height=250,
                    show_download_button=False
                )
            
            # KORREKT: coordinate_inputs in der richtigen Reihenfolge (x1, y1, x2, y2)
            coordinate_inputs = [img_input, bbox_x1, bbox_y1, bbox_x2, bbox_y2, mode_radio]
            
            img_input.upload(
                fn=mock_process_image_upload,
                inputs=[img_input],
                outputs=[preview_output, bbox_x1, bbox_y1, bbox_x2, bbox_y2]
            ).then(
                fn=mock_update_slider_for_image,
                inputs=[img_input],
                outputs=[bbox_x1, bbox_y1, bbox_x2, bbox_y2]
            )
            
            for slider in [bbox_x1, bbox_y1, bbox_x2, bbox_y2]:
                slider.release(
                    fn=mock_update_live_preview,
                    inputs=coordinate_inputs,
                    outputs=preview_output
                )
            
            mode_radio.change(
                fn=mock_update_info,
                inputs=[mode_radio],
                outputs=[pos_info, neg_info]
            ).then(
                fn=mock_update_live_preview,
                inputs=coordinate_inputs,
                outputs=preview_output
            )
  
            
            transform_btn.click(
                fn=mock_img_to_image,
                inputs=[
                    img_input, img_prompt, img_neg_prompt, 
                    strength_slider, img_steps, img_guidance, 
                    mode_radio, bbox_x1, bbox_y1, bbox_x2, bbox_y2
                ],
                outputs=[img_output, sam_raw_mask_output, processed_mask_output, 
                        pose_map_output, canny_map_output],
                concurrency_limit=1
            )
    
    return demo

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