app.py

"""
ImageForge — Custom AI image generation wrapper
Tab 1: Text-to-Image with character presets (fully working)
Tabs 2-6: stubbed placeholders, built out in later passes

Deploy target: Hugging Face Space, Gradio SDK, ZeroGPU hardware
"""

import gradio as gr
import spaces
import torch
import json
import os
import random
from datetime import datetime
from pathlib import Path

# ---------------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------------

PRESETS_FILE = "presets.json"
MODELS_FILE = "models.json"
OUTPUT_DIR = Path("generations")
OUTPUT_DIR.mkdir(exist_ok=True)

# Starting list of models — NOT the full ceiling. New open-source
# text-to-image models get released constantly (this list will go stale).
# Any Hugging Face model id can be added from the UI below without
# touching this code — see load_models()/add_model().
DEFAULT_MODEL_OPTIONS = {
    "Z-Image Turbo (fast, photorealistic)": "Tongyi-MAI/Z-Image-Turbo",
    "Qwen-Image (best prompt following + text)": "Qwen/Qwen-Image",
    "Flux Schnell (fast, Apache 2.0)": "black-forest-labs/FLUX.1-schnell",
    "Flux Krea (stronger realism)": "black-forest-labs/FLUX.1-Krea-dev",
}


def load_models():
    """Returns the combined model list: built-in defaults + anything
    the user has added via the UI, persisted in models.json."""
    if os.path.exists(MODELS_FILE):
        with open(MODELS_FILE, "r") as f:
            custom = json.load(f)
    else:
        custom = {}
    combined = dict(DEFAULT_MODEL_OPTIONS)
    combined.update(custom)
    return combined


def save_custom_models(custom):
    with open(MODELS_FILE, "w") as f:
        json.dump(custom, f, indent=2)


def add_model(display_name, model_id):
    """Add a new model to the dropdown by pasting any Hugging Face model id,
    e.g. 'THUDM/CogView4-6B' or a brand new release. No code edit needed.
    NOTE: only diffusers-compatible text-to-image pipelines will work here —
    if a new model needs custom loading code, it may need a small code update."""
    if not display_name.strip() or not model_id.strip():
        return gr.update(), "Both a display name and a model ID are required."
    if os.path.exists(MODELS_FILE):
        with open(MODELS_FILE, "r") as f:
            custom = json.load(f)
    else:
        custom = {}
    custom[display_name.strip()] = model_id.strip()
    save_custom_models(custom)
    return gr.update(choices=list(load_models().keys())), f"Added model '{display_name.strip()}'."


MODEL_OPTIONS = load_models()

# Default negative prompt applied unless user overrides it.
DEFAULT_NEGATIVE_PROMPT = (
    "blurry, distorted, deformed hands, extra limbs, extra fingers, "
    "watermark, text artifacts, low quality, off-model, inconsistent style"
)

# Quality presets -> (steps, guidance_scale)
QUALITY_PRESETS = {
    "Draft (fast)": {"steps": 8, "guidance": 2.5},
    "Final Quality (slower, better)": {"steps": 30, "guidance": 4.5},
}

_loaded_pipelines = {}  # cache so we don't reload a model every call


# ---------------------------------------------------------------------------
# PRESET STORAGE (simple JSON file)
# ---------------------------------------------------------------------------

def load_presets():
    if os.path.exists(PRESETS_FILE):
        with open(PRESETS_FILE, "r") as f:
            return json.load(f)
    return {}


def save_presets(presets):
    with open(PRESETS_FILE, "w") as f:
        json.dump(presets, f, indent=2)


def preset_names():
    return list(load_presets().keys())


def add_or_update_preset(name, description, style):
    if not name or not name.strip():
        return gr.update(), "Preset name can't be empty."
    presets = load_presets()
    presets[name.strip()] = {
        "description": description.strip(),
        "style": style.strip(),
    }
    save_presets(presets)
    return gr.update(choices=preset_names(), value=name.strip()), f"Saved preset '{name.strip()}'."


def delete_preset(name):
    presets = load_presets()
    if name in presets:
        del presets[name]
        save_presets(presets)
        return gr.update(choices=preset_names(), value=None), f"Deleted preset '{name}'."
    return gr.update(), "Nothing to delete."


def load_preset_into_fields(name):
    presets = load_presets()
    if name and name in presets:
        p = presets[name]
        return p["description"], p["style"]
    return "", ""


# ---------------------------------------------------------------------------
# PROMPT ENRICHMENT (hidden pipeline layer — point 1/2 from the spec)
# ---------------------------------------------------------------------------

def build_final_prompt(user_prompt, preset_name, style_override):
    """Silently combine: user prompt + character preset + style, so the
    person typing never has to re-type the character description every time."""
    parts = [user_prompt.strip()]

    presets = load_presets()
    if preset_name and preset_name in presets:
        p = presets[preset_name]
        if p.get("description"):
            parts.append(p["description"])
        if p.get("style") and not style_override:
            parts.append(p["style"])

    if style_override and style_override.strip():
        parts.append(style_override.strip())

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


# ---------------------------------------------------------------------------
# MODEL LOADING + GENERATION
# ---------------------------------------------------------------------------

def get_pipeline(model_key):
    model_id = load_models()[model_key]
    if model_id not in _loaded_pipelines:
        from diffusers import DiffusionPipeline
        pipe = DiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
        )
        pipe = pipe.to("cuda")
        _loaded_pipelines[model_id] = pipe
    return _loaded_pipelines[model_id]


@spaces.GPU(duration=90)
def generate_images(
    prompt,
    negative_prompt,
    model_key,
    quality_key,
    preset_name,
    style_override,
    seed,
    num_variants,
):
    settings = QUALITY_PRESETS[quality_key]
    final_prompt = build_final_prompt(prompt, preset_name, style_override)
    neg_prompt = negative_prompt.strip() if negative_prompt.strip() else DEFAULT_NEGATIVE_PROMPT

    pipe = get_pipeline(model_key)

    # seed handling: -1 means "random, but tell me what it was"
    base_seed = random.randint(0, 2**31 - 1) if seed is None or seed < 0 else int(seed)

    images = []
    used_seeds = []
    for i in range(int(num_variants)):
        this_seed = base_seed + i
        generator = torch.Generator(device="cuda").manual_seed(this_seed)
        result = pipe(
            prompt=final_prompt,
            negative_prompt=neg_prompt,
            num_inference_steps=settings["steps"],
            guidance_scale=settings["guidance"],
            generator=generator,
        )
        img = result.images[0]
        images.append(img)
        used_seeds.append(this_seed)

        # Auto-save every generation immediately — safety net so nothing
        # is ever lost even if the browser tab closes before download.
        char_tag = preset_name.replace(" ", "") if preset_name else "NoPreset"
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        filename = f"{char_tag}_seed{this_seed}_{timestamp}.png"
        img.save(OUTPUT_DIR / filename)

    seed_report = ", ".join(str(s) for s in used_seeds)
    status = f"Generated {len(images)} variant(s). Seeds used: {seed_report}\nPrompt sent: {final_prompt}"
    return images, status


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="ImageForge") as demo:
    gr.Markdown("# 🎨 ImageForge\nCustom AI image generation — built for the Bojo project, usable for any character/scene/asset work.")

    with gr.Tabs():
        # -------------------------------------------------------------
        # TAB 1: TEXT TO IMAGE
        # -------------------------------------------------------------
        with gr.Tab("1. Text to Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    prompt_box = gr.Textbox(
                        label="Describe the scene",
                        placeholder="Bojo standing in a sunlit meadow, looking surprised",
                        lines=3,
                    )
                    model_dropdown = gr.Dropdown(
                        label="Model",
                        choices=list(MODEL_OPTIONS.keys()),
                        value=list(MODEL_OPTIONS.keys())[0],
                    )
                    with gr.Accordion("Add a new model (any Hugging Face model ID)", open=False):
                        gr.Markdown(
                            "New open-source image models come out all the time — this list "
                            "isn't fixed. Paste any diffusers-compatible model ID from "
                            "huggingface.co to add it here permanently."
                        )
                        new_model_name_box = gr.Textbox(label="Display name", placeholder="e.g. CogView4")
                        new_model_id_box = gr.Textbox(
                            label="Hugging Face model ID", placeholder="e.g. THUDM/CogView4-6B"
                        )
                        add_model_btn = gr.Button("Add model")
                        add_model_status = gr.Textbox(label="Status", interactive=False)
                    quality_dropdown = gr.Dropdown(
                        label="Quality",
                        choices=list(QUALITY_PRESETS.keys()),
                        value="Draft (fast)",
                    )

                    gr.Markdown("### Character preset")
                    preset_dropdown = gr.Dropdown(
                        label="Load a saved character",
                        choices=preset_names(),
                        value=None,
                    )
                    style_override_box = gr.Textbox(
                        label="Style override (optional, replaces preset style just for this run)",
                        placeholder="e.g. watercolor, painterly, storybook illustration",
                    )

                    with gr.Accordion("Advanced (seed, negative prompt, variants)", open=False):
                        seed_box = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)
                        num_variants_box = gr.Slider(
                            label="Number of variants", minimum=1, maximum=4, step=1, value=2
                        )
                        negative_prompt_box = gr.Textbox(
                            label="Negative prompt (leave blank to use the default)",
                            placeholder=DEFAULT_NEGATIVE_PROMPT,
                            lines=2,
                        )

                    generate_btn = gr.Button("Generate", variant="primary")

                with gr.Column(scale=1):
                    gallery = gr.Gallery(label="Results", columns=2, height=500)
                    status_box = gr.Textbox(label="Status / prompt actually sent", lines=4, interactive=False)
                    gr.Markdown(
                        "All generations auto-save to this Space's storage as a backup. "
                        "Click any image in the gallery, then use the download icon to save it to your PC."
                    )

            generate_btn.click(
                fn=generate_images,
                inputs=[
                    prompt_box,
                    negative_prompt_box,
                    model_dropdown,
                    quality_dropdown,
                    preset_dropdown,
                    style_override_box,
                    seed_box,
                    num_variants_box,
                ],
                outputs=[gallery, status_box],
            )

            add_model_btn.click(
                fn=add_model,
                inputs=[new_model_name_box, new_model_id_box],
                outputs=[model_dropdown, add_model_status],
            )

            gr.Markdown("---")
            gr.Markdown("### Manage character presets")
            with gr.Row():
                with gr.Column():
                    preset_name_box = gr.Textbox(label="Preset name", placeholder="Bojo")
                    preset_desc_box = gr.Textbox(
                        label="Character description",
                        placeholder="a grey donkey with a patched blue vest, big expressive eyes, floppy ears",
                        lines=3,
                    )
                    preset_style_box = gr.Textbox(
                        label="Default style for this character",
                        placeholder="flat cartoon, thick outlines, warm color palette",
                        lines=2,
                    )
                    with gr.Row():
                        save_preset_btn = gr.Button("Save / Update preset")
                        delete_preset_btn = gr.Button("Delete selected preset")
                    preset_status_box = gr.Textbox(label="Preset status", interactive=False)

            # wire up preset management
            preset_dropdown.change(
                fn=load_preset_into_fields,
                inputs=[preset_dropdown],
                outputs=[preset_desc_box, preset_style_box],
            )
            save_preset_btn.click(
                fn=add_or_update_preset,
                inputs=[preset_name_box, preset_desc_box, preset_style_box],
                outputs=[preset_dropdown, preset_status_box],
            )
            delete_preset_btn.click(
                fn=delete_preset,
                inputs=[preset_dropdown],
                outputs=[preset_dropdown, preset_status_box],
            )

        # -------------------------------------------------------------
        # STUBS — built out in the next pass
        # -------------------------------------------------------------
        with gr.Tab("2. Image to Image (coming next)"):
            gr.Markdown("Reference-based editing with Qwen-Image-Edit / FLUX Kontext. Built after Tab 1 is confirmed working.")

        with gr.Tab("3. Reference & Consistency (coming next)"):
            gr.Markdown("Reference image lock + seed reuse + style presets, layered on top of Tab 1/2.")

        with gr.Tab("4. Character Sheet Builder (coming next)"):
            gr.Markdown("One reference run through multiple angle/pose presets in a batch.")

        with gr.Tab("5. Storyboard (coming next)"):
            gr.Markdown("Grid view of all scenes, click-to-regenerate, export as one board.")

        with gr.Tab("6. Image Editor (coming next)"):
            gr.Markdown("Crop, resize, mask-based inpainting on any saved generation.")

if __name__ == "__main__":
    demo.queue()
    demo.launch()