import gradio as gr import numpy as np import random import torch from diffusers import DiffusionPipeline import warnings import os from datetime import datetime import uuid # 경고 메시지 숨기기 warnings.filterwarnings('ignore', category=UserWarning) # 저장 디렉토리 생성 SAVE_DIR = "saved_images" if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR, exist_ok=True) # 장치 및 dtype 설정 dtype = torch.float32 if torch.cuda.is_available() else torch.float32 device = "cuda" if torch.cuda.is_available() else "cpu" # 모델 로드 pipe = DiffusionPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype, device_map="auto", use_safetensors=True ).to(device) # 메모리 최적화 pipe.enable_attention_slicing() if device == "cpu": pipe.enable_sequential_cpu_offload() MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 2048 def generate_diagram(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4): """FLUX AI를 사용하여 다이어그램 생성""" try: if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) with torch.no_grad(): image = pipe( prompt=prompt, width=width, height=height, num_inference_steps=num_inference_steps, generator=generator, guidance_scale=0.0 ).images[0] # 이미지 저장 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") unique_id = str(uuid.uuid4())[:8] filename = f"diagram_{timestamp}_{unique_id}.png" save_path = os.path.join(SAVE_DIR, filename) image.save(save_path) return image, seed except Exception as e: raise gr.Error(f"다이어그램 생성 중 오류 발생: {str(e)}") finally: if torch.cuda.is_available(): torch.cuda.empty_cache() # Enhanced examples with more detailed prompts and specific styling EXAMPLES = [ { "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": "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": "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": "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": "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": "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": "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": "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": "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": "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 } ] # Convert examples to Gradio format GRADIO_EXAMPLES = [ [example["prompt"], example["width"], example["height"]] for example in EXAMPLES ] def generate_diagram(prompt, width=1024, height=1024): """FLUX AI를 사용하여 다이어그램 생성""" try: # 시드 설정 seed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed) # 이미지 생성 image = pipeline( prompt=prompt, width=width, height=height, num_inference_steps=4, generator=generator, ).images[0] # 이미지 저장 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") unique_id = str(uuid.uuid4())[:8] filename = f"diagram_{timestamp}_{unique_id}.png" save_path = os.path.join(SAVE_DIR, filename) image.save(save_path) return image except Exception as e: raise gr.Error(f"다이어그램 생성 중 오류 발생: {str(e)}") # CSS 스타일 css=""" #col-container { margin: 0 auto; max-width: 520px; } """ # Gradio 인터페이스 생성 with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("""# FLUX 다이어그램 생성기 FLUX AI를 사용하여 아름다운 손그림 스타일의 다이어그램을 생성합니다 """) with gr.Row(): prompt = gr.Text( label="프롬프트", show_label=False, max_lines=1, placeholder="다이어그램 구조를 입력하세요...", container=False, ) run_button = gr.Button("생성", scale=0) result = gr.Image(label="생성된 다이어그램", show_label=False) with gr.Accordion("고급 설정", open=False): seed = gr.Slider( label="시드", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="랜덤 시드", value=True) with gr.Row(): width = gr.Slider( label="너비", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, ) height = gr.Slider( label="높이", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, ) num_inference_steps = gr.Slider( label="추론 단계 수", minimum=1, maximum=50, step=1, value=4, ) # 예제 추가 gr.Examples( examples=EXAMPLES, # 이전에 정의된 예제들 사용 fn=generate_diagram, inputs=[prompt], outputs=[result, seed], cache_examples=True ) # 이벤트 핸들러 gr.on( triggers=[run_button.click, prompt.submit], fn=generate_diagram, inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps], outputs=[result, seed] ) # 앱 실행 if __name__ == "__main__": demo.queue() demo.launch( server_name="0.0.0.0", server_port=7860, share=False, show_error=True, debug=True )