Update app.py

app.py

@@ -1,508 +1,311 @@
-# app.py
-import os
-import base64
-import streamlit as st
+import gradio as gr
 from gradio_client import Client
+import os
 from dotenv import load_dotenv
-from pathlib import Path
-import json
-import hashlib
-import time
-from typing import Dict, Any
 
 # Load environment variables
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Cache directory setup
-CACHE_DIR = Path("./cache")
-CACHE_DIR.mkdir(exist_ok=True)
-
-# Cached example diagrams
-CACHED_EXAMPLES = {
-    "literacy_mental": {
-        "title": "Literacy Mental Map",
-        "prompt": """A handrawn colorful mind map diagram, rugosity drawn lines, clear shapes, brain silhouette, text areas. must include the texts 
-        LITERACY/MENTAL
-        ├── PEACE [Dove Icon]
-        ├── HEALTH [Vitruvian Man ~60px]
-        ├── CONNECT [Brain-Mind Connection Icon]
-        ├── INTELLIGENCE
-        │   └── EVERYTHING [Globe Icon ~50px]
-        └── MEMORY
-            ├── READING [Book Icon ~40px]
-            ├── SPEED [Speedometer Icon]
-            └── CREATIVITY
-                └── INTELLIGENCE [Lightbulb + Infinity ~30px]""",
-        "width": 1024,
-        "height": 1024,
-        "seed": 1872187377,
-        "cache_path": "literacy_mental.png"
-    }
-}
-
-DIAGRAM_EXAMPLES = [
-    {
-        "title": "Project Management Flow",
-        "prompt": """A handrawn colorful mind map diagram, rugosity drawn lines, clear shapes, project management flow.
-        PROJECT MANAGEMENT
-        ├── INITIATION [Rocket Icon]
-        ├── PLANNING [Calendar Icon]
-        ├── EXECUTION [Gear Icon]
-        ├── MONITORING
-        │   └── CONTROL [Dashboard Icon]
-        └── CLOSURE [Checkmark Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Digital Marketing Strategy",
-        "prompt": """A handrawn colorful mind map diagram, rugosity drawn lines, modern style, marketing concept.
-        DIGITAL MARKETING
-        ├── SEO [Magnifying Glass]
-        ├── SOCIAL MEDIA [Network Icon]
-        ├── CONTENT
-        │   ├── BLOG [Document Icon]
-        │   └── VIDEO [Play Button]
-        └── ANALYTICS [Graph Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Software Development Lifecycle",
-        "prompt": """A handrawn colorful mind map diagram, technical style, software development flow.
-        SDLC
-        ├── REQUIREMENTS [Document Icon]
-        ├── DESIGN [Blueprint Icon]
-        ├── DEVELOPMENT [Code Icon]
-        ├── TESTING
-        │   ├── UNIT [Checkbox Icon]
-        │   └── INTEGRATION [Puzzle Icon]
-        └── DEPLOYMENT [Rocket Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Environmental Sustainability",
-        "prompt": """A handrawn colorful mind map diagram, nature-inspired style, environmental concept.
-        ECO-FRIENDLY
-        ├── REDUCE [Minus Icon]
-        ├── REUSE [Cycle Icon]
-        ├── RECYCLE [Triangle Icon]
-        └── ENERGY
-            ├── SOLAR [Sun Icon]
-            └── WIND [Windmill Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
+# Enhanced examples with more detailed prompts and specific styling
+EXAMPLES = [
     {
-        "title": "Business Strategy",
-        "prompt": """A handrawn colorful mind map diagram, professional style, business planning.
-        STRATEGY
-        ├── VISION [Eye Icon]
-        ├── MISSION [Target Icon]
-        ├── GOALS
-        │   ├── SHORT-TERM [Clock Icon]
-        │   └── LONG-TERM [Calendar Icon]
-        └── METRICS [Graph Icon]""",
+        "title": "Knowledge Tree",
+        "prompt": """A handrawn colorful mind map diagram, educational style, vibrant colors, clear hierarchy, golden ratio layout.
+        KNOWLEDGE
+        ├── ACQUISITION [Brain with Lightning ~60px]
+        │   ├── READING [Open Book with Glow]
+        │   ├── PRACTICE [Hands-on Tools]
+        │   └── OBSERVATION [Eye with Magnifier]
+        ├── PROCESSING [Gear Network ~50px]
+        │   ├── ANALYSIS [Graph Trending Up]
+        │   └── SYNTHESIS [Puzzle Pieces]
+        ├── RETENTION [Memory Chip ~45px]
+        │   ├── SHORT-TERM [Quick Flash]
+        │   └── LONG-TERM [Solid Archive]
+        └── APPLICATION
+            ├── CREATION [Artist Palette]
+            └── INNOVATION [Lightbulb Constellation]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Health & Wellness",
-        "prompt": """A handrawn colorful mind map diagram, wellness-focused style, health aspects.
-        WELLNESS
-        ├── PHYSICAL [Dumbbell Icon]
-        ├── MENTAL [Brain Icon]
-        ├── NUTRITION [Apple Icon]
-        └── SLEEP
-            ├── QUALITY [Star Icon]
-            └── DURATION [Clock Icon]""",
+        "title": "Digital Transformation",
+        "prompt": """A handrawn colorful mind map diagram, tech-focused style, neon accents, circuit board patterns.
+        DIGITAL TRANSFORM
+        ├── CLOUD [Cloud with Data Rain ~55px]
+        │   ├── STORAGE [Database Cluster]
+        │   └── COMPUTING [Server Array]
+        ├── AUTOMATION [Robot Arm ~50px]
+        │   ├── WORKFLOWS [Flowchart]
+        │   └── AI/ML [Neural Network]
+        ├── SECURITY [Shield Matrix ~45px]
+        │   ├── ENCRYPTION [Lock Code]
+        │   └── MONITORING [Radar Screen]
+        └── INTEGRATION
+            ├── APIS [Puzzle Connect]
+            └── MICROSERVICES [Building Blocks]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Innovation Process",
-        "prompt": """A handrawn colorful mind map diagram, creative style, innovation flow.
-        INNOVATION
-        ├── IDEATION [Lightbulb Icon]
-        ├── RESEARCH [Magnifier Icon]
-        ├── PROTOTYPING [Tools Icon]
-        └── TESTING
-            ├── FEEDBACK [Chat Icon]
-            └── ITERATION [Cycle Icon]""",
+        "title": "Creative Process",
+        "prompt": """A handrawn colorful mind map diagram, artistic style, watercolor effects, flowing connections.
+        CREATIVITY
+        ├── INSPIRATION [Constellation Stars ~60px]
+        │   ├── NATURE [Organic Patterns]
+        │   └── CULTURE [Global Icons]
+        ├── IDEATION [Floating Bubbles ~50px]
+        │   ├── BRAINSTORM [Thunder Cloud]
+        │   └── REFINEMENT [Diamond Polish]
+        ├── EXECUTION [Artist Tools ~45px]
+        │   ├── TECHNIQUE [Skilled Hands]
+        │   └── MEDIUM [Palette Mix]
+        └── PRESENTATION
+            ├── GALLERY [Frame Display]
+            └── FEEDBACK [Echo Ripples]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Data Science Workflow",
-        "prompt": """A handrawn colorful mind map diagram, technical style, data science process.
-        DATA SCIENCE
-        ├── COLLECTION [Database Icon]
-        ├── CLEANING [Filter Icon]
-        ├── ANALYSIS [Graph Icon]
-        └── VISUALIZATION
-            ├── CHARTS [Bar Chart Icon]
-            └── REPORTS [Document Icon]""",
+        "title": "Future Cities",
+        "prompt": """A handrawn colorful mind map diagram, futuristic style, holographic elements, sustainable themes.
+        SMART CITY
+        ├── MOBILITY [Hover Transport ~60px]
+        │   ├── AUTONOMOUS [Self-Driving]
+        │   └── CONNECTED [Network Grid]
+        ├── ENERGY [Solar Crystal ~55px]
+        │   ├── RENEWABLE [Green Power]
+        │   └── STORAGE [Battery Hub]
+        ├── LIVING [Eco Building ~50px]
+        │   ├── VERTICAL [Sky Gardens]
+        │   └── COMMUNITY [People Connect]
+        └── INFRASTRUCTURE
+            ├── AI GRID [Neural City]
+            └── ECO SYSTEM [Nature Tech]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Product Development",
-        "prompt": """A handrawn colorful mind map diagram, product-focused style, development stages.
-        PRODUCT
-        ├── RESEARCH [Magnifier Icon]
-        ├── DESIGN [Pencil Icon]
-        ├── BUILD [Tools Icon]
-        └── LAUNCH
-            ├── TESTING [Checkbox Icon]
-            └── RELEASE [Rocket Icon]""",
+        "title": "Health Evolution",
+        "prompt": """A handrawn colorful mind map diagram, medical style, DNA helix patterns, wellness focus.
+        HEALTH 3.0
+        ├── PREVENTION [Shield DNA ~60px]
+        │   ├── LIFESTYLE [Activity Pulse]
+        │   └── MONITORING [Health Watch]
+        ├── TREATMENT [Caduceus Tech ~55px]
+        │   ├── PERSONALIZED [DNA Code]
+        │   └── REGENERATIVE [Cell Renew]
+        ├── ENHANCEMENT [Upgrade Spiral ~50px]
+        │   ├── COGNITIVE [Brain Boost]
+        │   └── PHYSICAL [Body Optimize]
+        └── INTEGRATION
+            ├── AI HEALTH [Smart Doctor]
+            └── COMMUNITY [Global Care]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Time Management",
-        "prompt": """A handrawn colorful mind map diagram, organizational style, time concepts.
-        TIME
-        ├── PLANNING [Calendar Icon]
-        ├── PRIORITIES [Star Icon]
-        ├── EXECUTION [Clock Icon]
-        └── REVIEW
-            ├── DAILY [Sun Icon]
-            └── WEEKLY [Moon Icon]""",
+        "title": "Space Exploration",
+        "prompt": """A handrawn colorful mind map diagram, cosmic style, star field background, planetary elements.
+        SPACE FRONTIER
+        ├── DISCOVERY [Telescope Array ~60px]
+        │   ├── MAPPING [Star Charts]
+        │   └── ANALYSIS [Data Stream]
+        ├── TRAVEL [Rocket Launch ~55px]
+        │   ├── PROPULSION [Energy Core]
+        │   └── NAVIGATION [Space Map]
+        ├── COLONIZATION [Dome City ~50px]
+        │   ├── HABITATS [Life Sphere]
+        │   └── RESOURCES [Mine Extract]
+        └── RESEARCH
+            ├── ASTROBIOLOGY [Life Search]
+            └── PHYSICS [Space Time]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Customer Journey",
-        "prompt": """A handrawn colorful mind map diagram, user-centric style, customer experience.
-        CUSTOMER
-        ├── AWARENESS [Eye Icon]
-        ├── CONSIDERATION [Think Icon]
-        ├── PURCHASE [Cart Icon]
-        └── LOYALTY
-            ├── SUPPORT [Headset Icon]
-            └── FEEDBACK [Star Icon]""",
+        "title": "Ocean Innovation",
+        "prompt": """A handrawn colorful mind map diagram, marine style, wave patterns, aqua themes.
+        OCEAN TECH
+        ├── EXPLORATION [Deep Submersible ~60px]
+        │   ├── MAPPING [Sonar Wave]
+        │   └── RESEARCH [Lab Bubble]
+        ├── CONSERVATION [Marine Life ~55px]
+        │   ├── PROTECTION [Reef Shield]
+        │   └── RESTORATION [Growth Core]
+        ├── HARVESTING [Sustainable Net ~50px]
+        │   ├── ENERGY [Wave Power]
+        │   └── RESOURCES [Bio Extract]
+        └── MONITORING
+            ├── AI SYSTEMS [Smart Sensors]
+            └── ECOLOGY [Life Web]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Financial Planning",
-        "prompt": """A handrawn colorful mind map diagram, finance-focused style, money management.
-        FINANCE
-        ├── INCOME [Dollar Icon]
-        ├── SAVINGS [Piggy Bank Icon]
-        ├── INVESTMENTS [Graph Icon]
-        └── EXPENSES
-            ├── FIXED [Lock Icon]
-            └── VARIABLE [Scale Icon]""",
+        "title": "Quantum Computing",
+        "prompt": """A handrawn colorful mind map diagram, quantum style, wave-particle duality, matrix patterns.
+        QUANTUM TECH
+        ├── COMPUTATION [Qubit Matrix ~60px]
+        │   ├── PROCESSING [Wave Function]
+        │   └── ALGORITHMS [Code Quantum]
+        ├── APPLICATIONS [Use Cases ~55px]
+        │   ├── SIMULATION [Model World]
+        │   └── OPTIMIZATION [Peak Find]
+        ├── INFRASTRUCTURE [Q-Hardware ~50px]
+        │   ├── CONTROL [Pulse Shape]
+        │   └── COOLING [Zero Point]
+        └── DEVELOPMENT
+            ├── SOFTWARE [Q-Code Web]
+            └── INTEGRATION [Classical Bridge]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Team Structure",
-        "prompt": """A handrawn colorful mind map diagram, organizational style, team hierarchy.
-        TEAM
-        ├── LEADERSHIP [Crown Icon]
-        ├── MANAGERS [Briefcase Icon]
-        ├── SPECIALISTS [Star Icon]
-        └── SUPPORT
-            ├── HR [Heart Icon]
-            └── IT [Gear Icon]""",
+        "title": "Bio Engineering",
+        "prompt": """A handrawn colorful mind map diagram, biological style, DNA patterns, organic flow.
+        BIOTECH
+        ├── GENETICS [DNA Helix ~60px]
+        │   ├── EDITING [CRISPR Tool]
+        │   └── SYNTHESIS [Gene Build]
+        ├── APPLICATIONS [Lab Array ~55px]
+        │   ├── MEDICINE [Heal Cell]
+        │   └── AGRICULTURE [Grow Plus]
+        ├── PLATFORMS [Bio Factory ~50px]
+        │   ├── SENSORS [Live Detect]
+        │   └── PROCESSORS [Cell Compute]
+        └── INTEGRATION
+            ├── AI BIOLOGY [Smart Life]
+            └── ECOSYSTEM [Nature Net]""",
         "width": 1024,
         "height": 1024
     },
     {
-        "title": "Learning Methods",
-        "prompt": """A handrawn colorful mind map diagram, educational style, learning approaches.
-        LEARNING
-        ├── VISUAL [Eye Icon]
-        ├── AUDITORY [Ear Icon]
-        ├── READING [Book Icon]
-        └── KINESTHETIC
-            ├── PRACTICE [Hand Icon]
-            └── EXPERIENCE [Star Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Content Creation",
-        "prompt": """A handrawn colorful mind map diagram, creative style, content types.
-        CONTENT
-        ├── TEXT [Document Icon]
-        ├── IMAGES [Camera Icon]
-        ├── VIDEO [Play Icon]
-        └── AUDIO
-            ├── PODCAST [Microphone Icon]
-            └── MUSIC [Note Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Risk Management",
-        "prompt": """A handrawn colorful mind map diagram, analytical style, risk assessment.
-        RISK
-        ├── IDENTIFY [Eye Icon]
-        ├── ANALYZE [Graph Icon]
-        ├── MITIGATE [Shield Icon]
-        └── MONITOR
-            ├── TRACK [Chart Icon]
-            └── REPORT [Document Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Quality Assurance",
-        "prompt": """A handrawn colorful mind map diagram, technical style, quality control.
-        QUALITY
-        ├── STANDARDS [Checkmark Icon]
-        ├── TESTING [Magnifier Icon]
-        ├── INSPECTION [Eye Icon]
-        └── IMPROVEMENT
-            ├── FEEDBACK [Chat Icon]
-            └── ACTION [Gear Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Communication Channels",
-        "prompt": """A handrawn colorful mind map diagram, network style, communication methods.
-        COMMUNICATION
-        ├── EMAIL [Envelope Icon]
-        ├── CHAT [Message Icon]
-        ├── PHONE [Phone Icon]
-        └── MEETINGS
-            ├── VIRTUAL [Screen Icon]
-            └── IN-PERSON [Group Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Career Development",
-        "prompt": """A handrawn colorful mind map diagram, professional style, career growth.
-        CAREER
-        ├── SKILLS [Tools Icon]
-        ├── EXPERIENCE [Clock Icon]
-        ├── NETWORK [Connect Icon]
-        └── GOALS
-            ├── SHORT [Flag Icon]
-            └── LONG [Mountain Icon]""",
-        "width": 1024,
-        "height": 1024
-    },
-    {
-        "title": "Problem Solving",
-        "prompt": """A handrawn colorful mind map diagram, analytical style, solution process.
-        PROBLEM
-        ├── IDENTIFY [Eye Icon]
-        ├── ANALYZE [Brain Icon]
-        ├── SOLVE [Key Icon]
-        └── VERIFY
-            ├── TEST [Checkbox Icon]
-            └── REVIEW [Magnifier Icon]""",
+        "title": "AI Evolution",
+        "prompt": """A handrawn colorful mind map diagram, neural network style, digital patterns, intelligence flow.
+        AI FUTURE
+        ├── COGNITION [Brain Network ~60px]
+        │   ├── LEARNING [Growth Path]
+        │   └── REASONING [Logic Tree]
+        ├── PERCEPTION [Sensor Array ~55px]
+        │   ├── VISION [Eye Matrix]
+        │   └── LANGUAGE [Word Web]
+        ├── INTERACTION [Connect Hub ~50px]
+        │   ├── HUMAN [Bridge Link]
+        │   └── MACHINE [Code Path]
+        └── EVOLUTION
+            ├── CONSCIOUSNESS [Mind Spark]
+            └── CREATIVITY [Art Core]""",
         "width": 1024,
         "height": 1024
     }
 ]
 
-
-
-
-
-
-
-
-class DiagramCache:
-    def __init__(self, cache_dir: Path):
-        self.cache_dir = cache_dir
-        self.cache_dir.mkdir(exist_ok=True)
-        self._load_cache()
-    
-    def _load_cache(self):
-        """Load existing cache entries"""
-        self.cache_index = {}
-        if (self.cache_dir / "cache_index.json").exists():
-            with open(self.cache_dir / "cache_index.json", "r") as f:
-                self.cache_index = json.load(f)
-    
-    def _save_cache_index(self):
-        """Save cache index to disk"""
-        with open(self.cache_dir / "cache_index.json", "w") as f:
-            json.dump(self.cache_index, f)
-    
-    def _get_cache_key(self, params: Dict[str, Any]) -> str:
-        """Generate cache key from parameters"""
-        param_str = json.dumps(params, sort_keys=True)
-        return hashlib.md5(param_str.encode()).hexdigest()
-    
-    def get(self, params: Dict[str, Any]) -> Path:
-        """Get cached result if exists"""
-        cache_key = self._get_cache_key(params)
-        cache_info = self.cache_index.get(cache_key)
-        if cache_info:
-            cache_path = self.cache_dir / cache_info["filename"]
-            if cache_path.exists():
-                return cache_path
-        return None
-    
-    def put(self, params: Dict[str, Any], result_path: Path):
-        """Store result in cache"""
-        cache_key = self._get_cache_key(params)
-        filename = f"{cache_key}{result_path.suffix}"
-        cache_path = self.cache_dir / filename
-        
-        # Copy result to cache
-        with open(result_path, "rb") as src, open(cache_path, "wb") as dst:
-            dst.write(src.read())
-        
-        # Update index
-        self.cache_index[cache_key] = {
-            "filename": filename,
-            "timestamp": time.time(),
-            "params": params
-        }
-        self._save_cache_index()
-
-
-# Initialize cache
-diagram_cache = DiagramCache(CACHE_DIR)
-
-@st.cache_data
-def generate_cached_example(example_id: str) -> str:
-    """Generate and cache example diagram"""
-    example = CACHED_EXAMPLES[example_id]
-    client = Client("black-forest-labs/FLUX.1-schnell")
-    
-    # Check cache first
-    cache_path = diagram_cache.get(example)
-    if cache_path:
-        with open(cache_path, "rb") as f:
-            return base64.b64encode(f.read()).decode()
-    
-    # Generate new image
-    result = client.predict(
-        prompt=example["prompt"],
-        seed=example["seed"],
-        randomize_seed=False,
-        width=example["width"],
-        height=example["height"],
-        num_inference_steps=4,
-        api_name="/infer"
-    )
-    
-    # Cache the result
-    diagram_cache.put(example, Path(result))
-    
-    with open(result, "rb") as f:
-        return base64.b64encode(f.read()).decode()
-
-def generate_diagram(prompt: str, width: int, height: int, seed: int = None) -> str:
-    """Generate a new diagram"""
-    client = Client("black-forest-labs/FLUX.1-schnell")
-    params = {
-        "prompt": prompt,
-        "seed": seed if seed else 1872187377,
-        "width": width,
-        "height": height
-    }
-    
-    # Check cache first
-    cache_path = diagram_cache.get(params)
-    if cache_path:
-        with open(cache_path, "rb") as f:
-            return base64.b64encode(f.read()).decode()
-    
-    # Generate new image
+def generate_diagram(prompt, width=1024, height=1024):
+    """Generate a diagram using FLUX AI"""
     try:
+        client = Client("black-forest-labs/FLUX.1-schnell", hf_token=HF_TOKEN)
         result = client.predict(
-            prompt=prompt,
-            seed=params["seed"],
-            randomize_seed=False,
-            width=width,
-            height=height,
-            num_inference_steps=4,
+            prompt,
+            1872187377,  # seed
+            False,       # randomize_seed
+            width,
+            height,
+            4,          # num_inference_steps
             api_name="/infer"
         )
-        
-        # Cache the result
-        diagram_cache.put(params, Path(result))
-        
-        with open(result, "rb") as f:
-            return base64.b64encode(f.read()).decode()
+        return result
     except Exception as e:
-        st.error(f"Error generating diagram: {str(e)}")
-        return None
+        raise gr.Error(f"Error generating diagram: {str(e)}")
 
+# Convert examples to Gradio format
+GRADIO_EXAMPLES = [
+    [example["prompt"], example["width"], example["height"]]
+    for example in EXAMPLES
+]
 
-def main():
-    st.set_page_config(page_title="FLUX Diagram Generator", layout="wide", initial_sidebar_state="expanded")
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_diagram,
+    inputs=[
+        gr.Textbox(
+            label="Diagram Prompt",
+            placeholder="Enter your diagram structure...",
+            lines=10
+        ),
+        gr.Slider(
+            label="Width",
+            minimum=512,
+            maximum=2048,
+            step=128,
+            value=1024
+        ),
+        gr.Slider(
+            label="Height",
+            minimum=512,
+            maximum=2048,
+            step=128,
+            value=1024
+        )
+    ],
+    outputs=gr.Image(type="filepath", label="Generated Diagram"),
+    title="🎨 FLUX Diagram Generator",
+    description="""Generate beautiful hand-drawn style diagrams using FLUX AI. 
+    Create intricate mind maps with custom themes and detailed iconography.""",
+    article="""
+    ### Tips for Better Results
+    - Use descriptive style hints (e.g., "vibrant colors", "flowing connections")
+    - Include specific icon sizes (e.g., "~60px")
+    - Add theme-specific patterns (e.g., "circuit board patterns", "wave patterns")
+    - Use consistent hierarchy and formatting
+    - Combine both functional and decorative elements
     
-    if "initialized" not in st.session_state:
-        st.session_state.initialized = True
-        st.session_state.generated_images = {}
-
-    st.title("🎨 FLUX Diagram Generator")
-    st.markdown("Generate beautiful hand-drawn style diagrams using FLUX AI")
-
-    # Initialize client
-    client = Client("black-forest-labs/FLUX.1-schnell", hf_token=HF_TOKEN)
-
-    # Sidebar with examples
-    st.sidebar.title("📚 Example Templates")
-    selected_example = st.sidebar.selectbox(
-        "Choose a template",
-        options=range(len(DIAGRAM_EXAMPLES)),
-        format_func=lambda x: DIAGRAM_EXAMPLES[x]["title"]
-    )
-
-    # Main content
-    col1, col2 = st.columns([2, 1])
-
-    with col1:
-        example = DIAGRAM_EXAMPLES[selected_example]
-        prompt = st.text_area("Diagram Prompt", value=example["prompt"], height=200)
-        
-        with st.expander("Advanced Configuration"):
-            width = st.number_input("Width", min_value=512, max_value=2048, value=1024, step=128)
-            height = st.number_input("Height", min_value=512, max_value=2048, value=1024, step=128)
-
-        if st.button("🎨 Generate Diagram"):
-            try:
-                with st.spinner("Generating diagram..."):
-                    result = client.predict(
-                        prompt,
-                        1872187377,  # seed
-                        False,       # randomize_seed
-                        width,
-                        height,
-                        4,          # num_inference_steps
-                        api_name="/infer"
-                    )
-                    st.session_state.generated_images[prompt] = result
-                    st.image(result, caption="Generated Diagram", use_column_width=True)
-            except Exception as e:
-                st.error(f"Error generating diagram: {str(e)}")
-
-    with col2:
-        st.subheader("Tips for Better Results")
-        st.markdown("""
-        - Use clear hierarchical structures
-        - Include icon descriptions in brackets
-        - Keep text concise and meaningful
-        - Use consistent formatting
-        """)
-
-        st.subheader("Template Structure")
-        st.code("""
-        MAIN TOPIC
-        ├── SUBTOPIC 1 [Icon]
-        ├── SUBTOPIC 2 [Icon]
-        └── SUBTOPIC 3
-            ├── DETAIL 1 [Icon]
-            └── DETAIL 2 [Icon]
-        """)
+    ### Template Structure
+    ```
+    MAIN TOPIC
+    ├── PRIMARY [Icon ~60px]
+    │   ├── SUB-A [Specific Icon]
+    │   └── SUB-B [Detailed Icon]
+    ├── SECONDARY [Icon ~55px]
+    │   ├── SUB-C [Themed Icon]
+    │   └── SUB-D [Custom Icon]
+    └── TERTIARY
+        ├── SUB-E [Dynamic Icon]
+        └── SUB-F [Unique Icon]
+    ```
+    """,
+    examples=GRADIO_EXAMPLES,
+    cache_examples=True,
+    theme=gr.themes.Soft(),
+    css="""
+        .gradio-container {
+            font-family: 'Inter', sans-serif;
+        }
+        .gr-prose {
+            max-width: 850px;
+            margin: auto;
+        }
+        .gr-prose code {
+            background-color: #f3f4f6;
+            padding: 0.2em 0.4em;
+            border-radius: 0.3em;
+            font-size: 0.95em;
+        }
+        .gr-prose h3 {
+            margin-top: 1.5em;
+            color: #2563eb;
+        }
+    """
+)
 
+# Launch app
 if __name__ == "__main__":
-    main()
-
-    
+    # Set HF space configs
+    demo.queue(max_size=15)  # Increased queue size
+    demo.launch(
+        share=False,
+        show_error=True,
+        server_name="0.0.0.0",
+        server_port=7860,
+    )