""" 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()