app.py

import gradio as gr
import torch
import random
import time
from datetime import datetime
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
import json

# --- Configuration & Model Loading ---
# Replicating the logic of your ComfyUI JSON:
# - CheckpointLoaderSimple -> runwayml/stable-diffusion-v1-5
# - KSampler -> EulerAncestralDiscreteScheduler
# - EmptyLatentImage -> 512x512

MODEL_ID = "runwayml/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load the pipeline (simulating the backend compute)
print(f"Loading model {MODEL_ID} on {DEVICE}...")
try:
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        safety_checker=None # Disable for speed/demo purposes
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    if DEVICE == "cuda":
        pipe.to(DEVICE)
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")
    pipe = None

# --- Core Logic Functions ---

def generate_image_with_memory(
    prompt, 
    negative_prompt, 
    steps, 
    cfg, 
    seed, 
    memory_db
):
    """
    Generates an image and updates the 'Long Term Memory' database.
    """
    if pipe is None:
        return None, memory_db, "Error: Model not loaded. Check GPU availability."

    start_time = time.time()
    
    # Handle Seed (0 = random)
    if seed == 0:
        seed = random.randint(0, 2147483647)
    
    generator = torch.Generator(device=DEVICE).manual_seed(seed)
    
    # Generate Image (Replicating KSampler + VAE Decode)
    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(cfg),
            height=512,
            width=512,
            generator=generator
        )
        image = result.images[0]
    except Exception as e:
        return None, memory_db, f"Generation Failed: {str(e)}"

    # Update Memory Database
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    new_entry = {
        "timestamp": timestamp,
        "prompt": prompt,
        "negative": negative_prompt,
        "seed": seed,
        "steps": steps,
        "cfg": cfg
    }
    
    # Prepend new entry to memory
    updated_memory = [new_entry] + memory_db
    
    processing_time = f"{time.time() - start_time:.2f}s"
    
    return image, updated_memory, f"Generated in {processing_time} | Seed: {seed}"

def recall_prompt(memory_db, evt: gr.SelectData):
    """
    Recalls a specific prompt from history based on user selection.
    """
    if memory_db is None or len(memory_db) == 0:
        return "", "", 0, 20, 7.0
    
    index = evt.index[0]
    selected = memory_db[index]
    
    return (
        selected["prompt"], 
        selected["negative"], 
        int(selected["seed"]), 
        int(selected["steps"]), 
        float(selected["cfg"])
    )

def export_memory(memory_db):
    """
    Exports the memory database to a downloadable JSON file.
    """
    if not memory_db:
        return None
    json_str = json.dumps(memory_db, indent=4)
    return json_str

# --- Gradio 6 UI Construction ---

with gr.Blocks() as demo:
    
    # Header
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🧠 AI-Native Memory Core</h1>
        <p>Automated Long Term Super Intelligent Human Recall | Persistent Database</p>
        <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #007bff; text-decoration: none;">Built with anycoder</a>
    </div>
    """)
    
    with gr.Row():
        # --- Sidebar: Controls & Memory Database ---
        with gr.Sidebar(width=350):
            gr.Markdown("### ⚙️ Generation Parameters")
            
            with gr.Accordion("Advanced Settings", open=True):
                steps = gr.Slider(minimum=1, maximum=150, value=20, step=1, label="Sampling Steps", info="Higher = more detail")
                cfg = gr.Slider(minimum=1, maximum=30, value=7.0, step=0.5, label="CFG Scale", info="Prompt adherence")
                seed = gr.Number(value=0, label="Seed (0 = Random)", precision=0)
            
            gr.Markdown("---")
            gr.Markdown("### 💾 Persistent Memory Database")
            gr.Markdown("Click a row to recall prompt settings.")
            
            memory_display = gr.Dataframe(
                label="Prompt History",
                headers=["Time", "Prompt", "Seed"],
                datatype=["str", "str", "number"],
                interactive=False,
                wrap=True,
                height=300
            )
            
            with gr.Row():
                export_btn = gr.Button("💾 Export Memory (JSON)", size="sm")
                download_file = gr.File(visible=False, label="Download Database")

            gr.Markdown("---")
            gr.Markdown("### 📱 App Access")
            # Simulating the APK request via PWA/Install instructions
            gr.Markdown("""
            <small>
            To install on Android:<br>
            1. Open this in Chrome<br>
            2. Tap Menu (⋮)<br>
            3. Tap "Add to Home Screen"
            </small>
            """)
            app_link_btn = gr.Button("📲 Open App View", link="/?view=standalone", variant="primary")

        # --- Main Area: Chat & Generation ---
        with gr.Column(scale=2):
            with gr.Group():
                prompt_input = gr.Textbox(
                    label="Positive Prompt", 
                    placeholder="Describe your vision...", 
                    lines=3,
                    show_copy_button=True
                )
                negative_input = gr.Textbox(
                    label="Negative Prompt", 
                    value="low quality, blurry, distorted, bad anatomy",
                    lines=2,
                    visible=False # Collapsed for clean UI initially
                )
                
                with gr.Row():
                    toggle_neg = gr.Button("Show/Hide Negative", size="sm")
                    generate_btn = gr.Button("✨ Generate & Remember", variant="primary", scale=2)
            
            # Output Area
            with gr.Row():
                with gr.Column():
                    result_image = gr.Image(label="Generated Artifact", type="pil", height=512)
                    status_text = gr.Textbox(label="System Status", interactive=False)

    # --- State Management ---
    # We use Gradio State to hold the "Memory Database" in the session
    memory_state = gr.State(value=[])

    # --- Event Listeners ---

    # 1. Toggle Negative Prompt
    toggle_neg.click(lambda: gr.Textbox(visible=True), None, negative_input)

    # 2. Generate Image & Update Memory
    generate_btn.click(
        fn=generate_image_with_memory,
        inputs=[prompt_input, negative_input, steps, cfg, seed, memory_state],
        outputs=[result_image, memory_state, status_text],
        api_visibility="public"
    )

    # 3. Update Memory Table Display whenever state changes
    # We format the list of dicts into a dataframe-friendly format
    def format_memory_for_display(db):
        if not db:
            return []
        return [[row["timestamp"], row["prompt"][:50] + "...", row["seed"]] for row in db]

    # We use a dummy trigger or just chain it. 
    # In Gradio 6, we can use the `change` event on the state itself or chain outputs.
    # Here we chain the output of generate to update the view directly.
    # Note: To keep it simple, we modify the generate function to return the formatted dataframe if needed,
    # or we can use a separate update trigger. Let's do it inline for efficiency.
    
    # Redefining generate to return formatted table for immediate display
    def generate_with_ui_update(prompt, neg, steps, cfg, seed, db):
        img, new_db, status = generate_image_with_memory(prompt, neg, steps, cfg, seed, db)
        formatted_table = format_memory_for_display(new_db)
        return img, new_db, status, formatted_table

    generate_btn.click(
        fn=generate_with_ui_update,
        inputs=[prompt_input, negative_input, steps, cfg, seed, memory_state],
        outputs=[result_image, memory_state, status_text, memory_display],
        api_visibility="public"
    )

    # 4. Recall from Memory (Click on Table)
    memory_display.select(
        fn=recall_prompt,
        inputs=[memory_state],
        outputs=[prompt_input, negative_input, seed, steps, cfg],
        api_visibility="private"
    )

    # 5. Export Memory
    export_btn.click(
        fn=export_memory,
        inputs=[memory_state],
        outputs=[download_file],
        api_visibility="private"
    )

# --- Launch Configuration ---
# Gradio 6 Syntax: Theme and parameters go in launch(), not Blocks()

demo.launch(
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="blue",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
        block_border_width="1px",
        block_border_color="*neutral_200"
    ),
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
        {"label": "Hugging Face", "url": "https://huggingface.co"}
    ],
    share=False # Set to True if you want a public link immediately
)