# 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()