app.py

import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
import spaces

# Basic settings
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model (FLUX.1-schnell)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=dtype
).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

@spaces.GPU()
def generate_image(prompt, seed, randomize_seed, width, height, steps, guidance_scale):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=steps,
        generator=generator,
        guidance_scale=guidance_scale
    ).images[0]
    return image, seed

def set_prompt(example_text):
    return example_text

# Example prompts per tab (each with "AI Electric Bicycle", "Dog-shaped Toy Robot", "Drone")
example_prompts = {
    "Product Design": [
        "A futuristic product design of an AI electric bicycle with modern aesthetics and advanced technology.",
        "A playful product design for a dog-shaped toy robot, cute and interactive.",
        "An innovative product design for a drone, sleek and highly functional."
    ],
    "Flowchart": [
        "A flowchart outlining the design and development process for an AI electric bicycle—from concept to production.",
        "A flowchart detailing the workflow for designing a dog-shaped toy robot, from brainstorming to testing.",
        "A flowchart mapping out the design process for a drone, including ideation and final assembly."
    ],
    "Infographic": [
        "An infographic presenting key statistics and design features of an AI electric bicycle, with modern icons and graphics.",
        "An infographic showing the design evolution of a dog-shaped toy robot, with sketches and key milestones.",
        "An infographic summarizing the innovative features of a drone design, with visual data and timelines."
    ],
    "Mockup": [
        "A mockup sketch of an AI electric bicycle interface and dashboard, showcasing futuristic control panels.",
        "A mockup wireframe for a dog-shaped toy robot's app interface, highlighting playful interactions.",
        "A mockup of a drone's control system and layout, featuring modern UX elements."
    ],
    "Diagram": [
        "A diagram illustrating the structural design and workflow of an AI electric bicycle, with clear stages and connections.",
        "A diagram outlining the components and design process of a dog-shaped toy robot, with labeled sections.",
        "A diagram showcasing the technical design and assembly process of a drone, with step-by-step visuals."
    ]
}

# CSS style: prevent example Accordion from expanding into the sidebar
css = """
* {
    box-sizing: border-box;
}
body {
    background: linear-gradient(135deg, #667eea, #764ba2);
    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
    color: #333;
    margin: 0;
    padding: 0;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 30px 40px;
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    margin: 40px auto;
    max-width: 1200px;
}
.sidebar {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
button, .btn {
    background: linear-gradient(90deg, #ff8a00, #e52e71);
    border: none;
    color: #fff;
    padding: 12px 24px;
    text-transform: uppercase;
    font-weight: bold;
    letter-spacing: 1px;
    border-radius: 5px;
    cursor: pointer;
    transition: transform 0.2s ease-in-out;
}
button:hover, .btn:hover {
    transform: scale(1.05);
}
.example-accordion {
    max-width: 100% !important;
}
"""

with gr.Blocks(css=css, title="Business Agent") as demo:
    gr.Markdown(
        """
        <div style="text-align:center;">
            <h1>Business Agent</h1>
            <p>Generate design concepts and workflow diagrams for your business needs using our multi-tab interface.</p>
            <p><strong>AI Spaces:</strong> Product Design, Flowchart, Infographic, Mockup, Diagram</p>
        </div>
        """
    )
    
    with gr.Row():
        # Left sidebar: Common generation parameters
        with gr.Sidebar(label="Parameters", open=True):
            gr.Markdown("### Generation Parameters")
            seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            width_slider = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            height_slider = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            steps_slider = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=20)
            guidance_slider = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=15.0, step=0.5, value=7.5)
        
        # Main area: Tabbed interface for different design types (first tab: Product Design)
        with gr.Column(scale=8):
            with gr.Tabs():
                # First tab: Product Design
                with gr.Tab("Product Design"):
                    pd_prompt = gr.Textbox(label="Product Design Prompt", placeholder="Enter a product design concept...", lines=5)
                    pd_generate = gr.Button("Generate Product Design")
                    pd_image = gr.Image(label="Generated Product Design")
                    with gr.Accordion("Example Prompts", open=False, elem_classes="example-accordion"):
                        for ex in example_prompts["Product Design"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=pd_prompt)
                
                # Second tab: Flowchart
                with gr.Tab("Flowchart"):
                    flow_prompt = gr.Textbox(label="Flowchart Prompt", placeholder="Enter a flowchart description...", lines=5)
                    flow_generate = gr.Button("Generate Flowchart")
                    flow_image = gr.Image(label="Generated Flowchart")
                    with gr.Accordion("Example Prompts", open=False, elem_classes="example-accordion"):
                        for ex in example_prompts["Flowchart"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=flow_prompt)
                
                # Third tab: Infographic
                with gr.Tab("Infographic"):
                    info_prompt = gr.Textbox(label="Infographic Prompt", placeholder="Enter an infographic description...", lines=5)
                    info_generate = gr.Button("Generate Infographic")
                    info_image = gr.Image(label="Generated Infographic")
                    with gr.Accordion("Example Prompts", open=False, elem_classes="example-accordion"):
                        for ex in example_prompts["Infographic"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=info_prompt)
                
                # Fourth tab: Mockup
                with gr.Tab("Mockup"):
                    mock_prompt = gr.Textbox(label="Mockup Prompt", placeholder="Enter a mockup description...", lines=5)
                    mock_generate = gr.Button("Generate Mockup")
                    mock_image = gr.Image(label="Generated Mockup")
                    with gr.Accordion("Example Prompts", open=False, elem_classes="example-accordion"):
                        for ex in example_prompts["Mockup"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=mock_prompt)
                
                # Fifth tab: Diagram
                with gr.Tab("Diagram"):
                    diag_prompt = gr.Textbox(label="Diagram Prompt", placeholder="Enter a diagram description...", lines=5)
                    diag_generate = gr.Button("Generate Diagram")
                    diag_image = gr.Image(label="Generated Diagram")
                    with gr.Accordion("Example Prompts", open=False, elem_classes="example-accordion"):
                        for ex in example_prompts["Diagram"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=diag_prompt)
    
    # Bind events for generation buttons
    pd_generate.click(
        fn=generate_image,
        inputs=[pd_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[pd_image, seed_slider]
    )
    
    flow_generate.click(
        fn=generate_image,
        inputs=[flow_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[flow_image, seed_slider]
    )
    
    info_generate.click(
        fn=generate_image,
        inputs=[info_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[info_image, seed_slider]
    )
    
    mock_generate.click(
        fn=generate_image,
        inputs=[mock_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[mock_image, seed_slider]
    )
    
    diag_generate.click(
        fn=generate_image,
        inputs=[diag_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[diag_image, seed_slider]
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )