Update app.py

app.py

@@ -1,229 +1,283 @@
-import random
+# app.py
 import os
-import uuid
-from datetime import datetime
-import gradio as gr
-import numpy as np
-import spaces
-import torch
-from diffusers import DiffusionPipeline
-from PIL import Image
-
-# Create permanent storage directory
-SAVE_DIR = "saved_images"  # Gradio will handle the persistence
-if not os.path.exists(SAVE_DIR):
-    os.makedirs(SAVE_DIR, exist_ok=True)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-repo_id = "black-forest-labs/FLUX.1-dev"
-adapter_id = "openfree/korea-president-yoon"
-
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-pipeline.load_lora_weights(adapter_id)
-pipeline = pipeline.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def save_generated_image(image, prompt):
-    # Generate unique filename with timestamp
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    unique_id = str(uuid.uuid4())[:8]
-    filename = f"{timestamp}_{unique_id}.png"
-    filepath = os.path.join(SAVE_DIR, filename)
+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"
+    }
+}
+
+# Example diagrams for various use cases
+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
+    }
+]
+
+# Add 15 more examples
+ADDITIONAL_EXAMPLES = [
+    {
+        "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
+    }
+    # ... (나머지 예제들)
+]
+
+
+
+
+
+
+
+
+
+
+
+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)
     
-    # Save the image
-    image.save(filepath)
+    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)
     
-    # Save metadata
-    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
-    with open(metadata_file, "a", encoding="utf-8") as f:
-        f.write(f"{filename}|{prompt}|{timestamp}\n")
+    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()
     
-    return filepath
+    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()
+
 
-def load_generated_images():
-    if not os.path.exists(SAVE_DIR):
-        return []
+# 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")
     
-    # Load all images from the directory
-    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
-                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
-    # Sort by creation time (newest first)
-    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
-    return image_files
-
-def load_predefined_images():
-    # Return empty list since we're not using predefined images
-    return []
-
-@spaces.GPU(duration=120)
-def inference(
-    prompt: str,
-    seed: int,
-    randomize_seed: bool,
-    width: int,
-    height: int,
-    guidance_scale: float,
-    num_inference_steps: int,
-    lora_scale: float,
-    progress: gr.Progress = gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
+    # 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()
     
-    image = pipeline(
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-    ).images[0]
+    # 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"
+    )
     
-    # Save the generated image
-    filepath = save_generated_image(image, prompt)
+    # Cache the result
+    diagram_cache.put(example, Path(result))
     
-    # Return the image, seed, and updated gallery
-    return image, seed, load_generated_images()
+    with open(result, "rb") as f:
+        return base64.b64encode(f.read()).decode()
 
-examples = [
-    "A man playing fetch with a golden retriever in a sunny park. He wears casual weekend clothes and throws a red frisbee with joy. The dog leaps gracefully through the air, tail wagging with excitement. Warm afternoon sunlight filters through the trees, creating a peaceful scene of companionship. [president yoon]",
+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
 
-    "A soldier standing at attention in full military gear, holding a standard-issue rifle. His uniform is crisp and properly adorned with medals. Behind him, other soldiers march in formation during a military parade. The scene conveys discipline and duty. [president yoon]",
 
-    "A medieval knight in gleaming armor, holding an ornate sword and shield. He stands proudly in front of a majestic castle, his cape flowing in the wind. The shield bears intricate heraldic designs, and sunlight glints off his polished armor. [president yoon]",
+def main():
+    st.set_page_config(page_title="FLUX Diagram Generator", layout="wide")
+    
+    st.title("🎨 FLUX Diagram Generator")
+    st.markdown("Generate beautiful hand-drawn style diagrams using FLUX AI")
 
-    "A charismatic political leader addressing a crowd from a podium. He wears a well-fitted suit and gestures confidently while speaking. The audience fills a large plaza, holding supportive banners and signs. News cameras capture the moment as he delivers his speech. [president yoon]",
+    # Sidebar for 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"]
+    )
 
-    "A man enjoying a peaceful morning at home, reading a newspaper at his breakfast table. He wears comfortable home clothes and sips coffee from a favorite mug. Sunlight streams through the kitchen window, and a house plant adds a touch of nature to the cozy domestic scene. [president yoon]",
+    # Main content area
+    col1, col2 = st.columns([2, 1])
 
-    "A businessman walking confidently through a modern office building. He carries a leather briefcase and wears a tailored navy suit. Floor-to-ceiling windows reveal a cityscape behind him, and his expression shows determination and purpose. [president yoon]"
-]
+    with col1:
+        # Input area
+        prompt = st.text_area(
+            "Diagram Prompt",
+            value=DIAGRAM_EXAMPLES[selected_example]["prompt"],
+            height=200
+        )
 
-css = """
-footer {
-    visibility: hidden;
-}
-"""
+        # Configuration
+        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)
+            seed = st.number_input("Seed (optional)", value=None, step=1)
 
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
-    gr.HTML('<div class="title"> President Yoon in KOREA </div>')
+        if st.button("🎨 Generate Diagram"):
+            with st.spinner("Generating your diagram..."):
+                result = generate_diagram(prompt, width, height, seed)
+                if result:
+                    st.image(result, caption="Generated Diagram", use_column_width=True)
 
-    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fopenfree-korea-president-yoon.hf.space">
-               <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fopenfree-korea-president-yoon.hf.space&countColor=%23263759" />
-               </a>""")
-    
-    with gr.Tabs() as tabs:
-        with gr.Tab("Generation"):
-            with gr.Column(elem_id="col-container"):
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt",
-                        container=False,
-                    )
-                    run_button = gr.Button("Run", scale=0)
-
-                result = gr.Image(label="Result", show_label=False)
-
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=42,
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=1024,
-                        )
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=768,
-                        )
-
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=0.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=3.5,
-                        )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=30,
-                        )
-                        lora_scale = gr.Slider(
-                            label="LoRA scale",
-                            minimum=0.0,
-                            maximum=1.0,
-                            step=0.1,
-                            value=1.0,
-                        )
-
-                gr.Examples(
-                    examples=examples,
-                    inputs=[prompt],
-                    outputs=[result, seed],
-                )
-
-        with gr.Tab("Gallery"):
-            gallery_header = gr.Markdown("### Generated Images Gallery")
-            generated_gallery = gr.Gallery(
-                label="Generated Images",
-                columns=6,
-                show_label=False,
-                value=load_generated_images(),
-                elem_id="generated_gallery",
-                height="auto"
-            )
-            refresh_btn = gr.Button("🔄 Refresh Gallery")
-
-
-    # Event handlers
-    def refresh_gallery():
-        return load_generated_images()
-
-    refresh_btn.click(
-        fn=refresh_gallery,
-        inputs=None,
-        outputs=generated_gallery,
-    )
+    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
+        """)
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=inference,
-        inputs=[
-            prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            lora_scale,
-        ],
-        outputs=[result, seed, generated_gallery],
-    )
+        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()
 
-demo.queue()
-demo.launch()
+