app.py

# app.py
import os
import base64
import streamlit as st
from gradio_client import Client
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
    },
    {
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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]""",
        "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
    try:
        result = client.predict(
            prompt=prompt,
            seed=params["seed"],
            randomize_seed=False,
            width=width,
            height=height,
            num_inference_steps=4,
            api_name="/infer"
        )
        
        # Cache the result
        diagram_cache.put(params, Path(result))
        
        with open(result, "rb") as f:
            return base64.b64encode(f.read()).decode()
    except Exception as e:
        st.error(f"Error generating diagram: {str(e)}")
        return None


def main():
    st.set_page_config(page_title="FLUX Diagram Generator", layout="wide", initial_sidebar_state="expanded")
    
    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]
        """)

if __name__ == "__main__":
    main()