Spaces:
Running
on
Zero
Running
on
Zero
# 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) | |
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() | |