import gradio as gr
from daggr import GradioNode, FnNode, Graph
import json
from datetime import datetime

# Node 1: Generate image using Z-Image-Turbo
z_image_turbo = GradioNode(
    space_or_url="Tongyi-MAI/Z-Image-Turbo",
    api_name="/generate",  # Common API name, might need adjustment
    inputs={
        "prompt": gr.Textbox(label="Prompt", placeholder="A beautiful landscape..."),
        "negative_prompt": gr.Textbox(label="Negative Prompt", value=""),
        "width": gr.Slider(512, 1024, value=1024, step=64, label="Width"),
        "height": gr.Slider(512, 1024, value=1024, step=64, label="Height"),
        "num_inference_steps": gr.Slider(1, 50, value=4, step=1, label="Steps"),  # Turbo models usually need few steps
        "guidance_scale": gr.Slider(1.0, 10.0, value=3.0, step=0.1, label="Guidance Scale"),
    },
    outputs={
        "image": gr.Image(label="Generated Image"),
        "info": gr.JSON(label="Generation Info")
    },
    postprocess=lambda original, target: {
        "image": original[0] if isinstance(original, list) else original,
        "info": {"timestamp": datetime.now().isoformat()}
    }
)

# Node 2: Optional metadata formatter
def format_metadata(prompt: str, width: int, height: int) -> str:
    return f"Prompt: {prompt}\nSize: {width}x{height}\nModel: Z-Image-Turbo"

metadata_node = FnNode(
    fn=format_metadata,
    inputs={
        "prompt": gr.Textbox(label="Prompt"),
        "width": gr.Number(label="Width"),
        "height": gr.Number(label="Height")
    },
    outputs={"metadata": gr.Textbox(label="Metadata", lines=3)}
)

# Create graph
graph = Graph(
    name="Z-Image-Turbo Generator",
    nodes=[z_image_turbo, metadata_node]
)

graph.launch()