Update app.py

app.py

@@ -2,50 +2,20 @@ import spaces
 import argparse
 import os
 import time
-import gc
 from os import path
-import shutil
-from datetime import datetime
-import traceback
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
-import gradio as gr
-import torch
-from diffusers import FluxPipeline
-from diffusers.pipelines.stable_diffusion import safety_checker
-from PIL import Image
 
-# Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
-# GPU 메모리 설정 최적화
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.benchmark = True  # 반복적인 크기의 입력에 대해 성능 향상
+import gradio as gr
+import torch
+from diffusers import FluxPipeline
 
-def filter_prompt(prompt):
-    # 부적절한 키워드 목록
-    inappropriate_keywords = [
-        # 음란/성적 키워드
-        "nude", "naked", "nsfw", "porn", "sex", "explicit", "adult", "xxx",
-        "erotic", "sensual", "seductive", "provocative", "intimate",
-        # 폭력적 키워드
-        "violence", "gore", "blood", "death", "kill", "murder", "torture",
-        # 기타 부적절한 키워드
-        "drug", "suicide", "abuse", "hate", "discrimination"
-    ]
-    
-    prompt_lower = prompt.lower()
-    
-    # 부적절한 키워드 체크
-    for keyword in inappropriate_keywords:
-        if keyword in prompt_lower:
-            return False, "부적절한 내용이 포함된 프롬프트입니다."
-            
-    return True, prompt
+torch.backends.cuda.matmul.allow_tf32 = True
 
 class timer:
     def __init__(self, method_name="timed process"):
@@ -57,320 +27,172 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
-# 글로벌 변수로 파이프라인 선언
-pipe = None
+if not path.exists(cache_path):
+    os.makedirs(cache_path, exist_ok=True)
 
-# 모델 초기화 함수 (지연 로딩)
-def initialize_model():
-    global pipe
-    
-    # 이미 로드된 경우 다시 로드하지 않음
-    if pipe is not None:
-        return True
-    
-    try:
-        if not path.exists(cache_path):
-            os.makedirs(cache_path, exist_ok=True)
-        
-        # 메모리 확보를 위한 가비지 컬렉션 실행
-        gc.collect()
-        torch.cuda.empty_cache()
-        
-        with timer("모델 로딩"):
-            pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-            lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
-            pipe.load_lora_weights(lora_path)
-            pipe.fuse_lora(lora_scale=0.125)
-            
-            # 주의: 여기서 device를 명시적으로 지정 (모든 컴포넌트에 적용)
-            pipe.to(device="cuda", dtype=torch.bfloat16)
-            
-            # 안전 검사기 추가 및 올바른 장치로 이동
-            pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
-            if hasattr(pipe, 'safety_checker') and pipe.safety_checker is not None:
-                pipe.safety_checker.to("cuda")
-            
-        print("모델 로딩 완료")
-        return True
-    except Exception as e:
-        print(f"모델 로딩 중 오류 발생: {str(e)}")
-        traceback.print_exc()
-        return False
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
+
+# Define example prompts
+example_prompts = [
+    "A cyberpunk cityscape at night with neon lights reflecting in puddles, towering skyscrapers and flying cars",
+    "An ethereal fairy with translucent iridescent wings standing in an enchanted forest with glowing mushrooms and floating light particles",
+    "A majestic dragon soaring through stormy clouds above jagged mountain peaks as lightning strikes in the background",
+    "A futuristic space station orbiting a vibrant nebula with multiple colorful ringed planets visible through a massive observation window",
+    "An underwater scene of an ancient lost city with ornate temples and statues covered in bioluminescent coral and swimming sea creatures"
+]
 
+# Custom CSS for neon theme
 css = """
-footer {display: none !important}
-.gradio-container {
-    max-width: 1200px;
-    margin: auto;
-}
-.contain {
-    background: rgba(255, 255, 255, 0.05);
-    border-radius: 12px;
-    padding: 20px;
-}
-.generate-btn {
-    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
-    border: none !important;
-    color: white !important;
-}
-.generate-btn:hover {
-    transform: translateY(-2px);
-    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
-}
-.title {
-    text-align: center;
-    font-size: 2.5em;
-    font-weight: bold;
-    margin-bottom: 1em;
-    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%);
-    -webkit-background-clip: text;
-    -webkit-text-fill-color: transparent;
-}
-.output-image {
-    width: 100% !important;
-    max-width: 100% !important;
-}
-.contain > div {
-    width: 100% !important;
-    max-width: 100% !important;
-}
-.fixed-width {
-    width: 100% !important;
-    max-width: 100% !important;
-}
-.loading-indicator {
-    text-align: center;
-    padding: 20px;
-    font-weight: bold;
-    color: #4B79A1;
-}
-.error-message {
-    background-color: rgba(255, 0, 0, 0.1);
-    color: red;
-    padding: 10px;
-    border-radius: 8px;
-    margin-top: 10px;
-    text-align: center;
-}
+    .neon-container {
+        background: linear-gradient(to right, #000428, #004e92);
+        border-radius: 16px;
+        box-shadow: 0 0 15px #00f3ff, 0 0 25px #00f3ff;
+    }
+    .neon-title {
+        text-shadow: 0 0 5px #fff, 0 0 10px #fff, 0 0 15px #0073e6, 0 0 20px #0073e6, 0 0 25px #0073e6;
+        color: #ffffff;
+        font-weight: bold !important;
+    }
+    .neon-text {
+        color: #00ff9d;
+        text-shadow: 0 0 5px #00ff9d;
+    }
+    .neon-button {
+        box-shadow: 0 0 5px #ff00dd, 0 0 10px #ff00dd !important;
+        background: linear-gradient(90deg, #ff00dd, #8b00ff) !important;
+        border: none !important;
+        color: white !important;
+        font-weight: bold !important;
+    }
+    .neon-button:hover {
+        box-shadow: 0 0 10px #ff00dd, 0 0 20px #ff00dd !important;
+    }
+    .neon-input {
+        border: 1px solid #00f3ff !important;
+        box-shadow: 0 0 5px #00f3ff !important;
+    }
+    .neon-slider > div {
+        background: linear-gradient(90deg, #00ff9d, #00f3ff) !important;
+    }
+    .neon-slider > div > div {
+        background: #ff00dd !important;
+        box-shadow: 0 0 5px #ff00dd !important;
+    }
+    .neon-card {
+        background-color: rgba(0, 0, 0, 0.7) !important;
+        border: 1px solid #00f3ff !important;
+        box-shadow: 0 0 10px #00f3ff !important;
+        padding: 16px !important;
+        border-radius: 8px !important;
+    }
+    .neon-example {
+        background: rgba(0, 0, 0, 0.5) !important;
+        border: 1px solid #00ff9d !important;
+        border-radius: 8px !important;
+        padding: 8px !important;
+        color: #00ff9d !important;
+        box-shadow: 0 0 5px #00ff9d !important;
+        margin: 4px !important;
+        cursor: pointer !important;
+    }
+    .neon-example:hover {
+        box-shadow: 0 0 10px #00ff9d, 0 0 15px #00ff9d !important;
+        background: rgba(0, 255, 157, 0.2) !important;
+    }
 """
 
-# Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    gr.HTML('<div class="title">AI Image Generator</div>')
-    gr.HTML('<div style="text-align: center; margin-bottom: 2em; color: #666;">Create stunning images from your descriptions</div>')
-    
-    gr.HTML("""
-        <div style="color: red; margin-bottom: 1em; text-align: center; padding: 10px; background: rgba(255,0,0,0.1); border-radius: 8px;">
-            ⚠️ Explicit or inappropriate content cannot be generated.
-        </div>
-    """)
+    with gr.Blocks(elem_classes=["neon-container"]):
+        gr.Markdown(
+            """
+            <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+                <h1 style="font-size: 3rem; font-weight: 700; margin-bottom: 1rem; display: contents;" class="neon-title">FLUX: Fast & Furious</h1>
+                <p style="font-size: 1.2rem; margin-bottom: 1.5rem;" class="neon-text">AutoML team from ByteDance</p>
+            </div>
+            """
+        )
 
-    # 상태 표시 변수 (HTML 대신 Textbox 사용)
-    error_message = gr.Textbox(
-        value="", 
-        label="Error", 
-        visible=False, 
-        elem_classes=["error-message"]
-    )
-    
-    loading_status = gr.Textbox(
-        value="", 
-        label="Status", 
-        visible=False, 
-        elem_classes=["loading-indicator"]
-    )
-
-    with gr.Row():
-        with gr.Column(scale=3):
-            prompt = gr.Textbox(
-                label="Image Description",
-                placeholder="Describe the image you want to create...",
-                lines=3
-            )
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=1152,
-                        step=64,
-                        value=1024
-                    )
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=256,
-                        maximum=1152,
-                        step=64,
-                        value=1024
+        with gr.Row():
+            with gr.Column(scale=3, elem_classes=["neon-card"]):
+                with gr.Group():
+                    prompt = gr.Textbox(
+                        label="Your Image Description",
+                        placeholder="E.g., A serene landscape with mountains and a lake at sunset",
+                        lines=3,
+                        elem_classes=["neon-input"]
                     )
-                
-                with gr.Row():
-                    steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=6,
-                        maximum=25,
-                        step=1,
-                        value=8
-                    )
-                    scales = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=0.0,
-                        maximum=5.0,
-                        step=0.1,
-                        value=3.5
-                    )
-                
-                def get_random_seed():
-                    return int(torch.randint(0, 1000000, (1,)).item())
-                
-                seed = gr.Number(
-                    label="Seed (random by default, set for reproducibility)",
-                    value=get_random_seed(),
-                    precision=0
-                )
-                
-                randomize_seed = gr.Button("🎲 Randomize Seed", elem_classes=["generate-btn"])
-            
-            generate_btn = gr.Button(
-                "✨ Generate Image",
-                elem_classes=["generate-btn"]
-            )
-            
-            gr.HTML("""
-                <div style="margin-top: 1em; padding: 1em; border-radius: 8px; background: rgba(255, 255, 255, 0.05);">
-                    <h4 style="margin: 0 0 0.5em 0;">Example Prompts:</h4>
-                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
-                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌅 Cinematic Landscape</p>
-                        <p style="margin: 0; font-style: italic;">"A breathtaking mountain vista at golden hour, dramatic sunbeams piercing through clouds, snow-capped peaks reflecting warm light, ultra-high detail photography, artistically composed, award-winning landscape photo, shot on Hasselblad"</p>
-                    </div>
-                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
-                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🖼️ Fantasy Portrait</p>
-                        <p style="margin: 0; font-style: italic;">"Ethereal portrait of an elven queen with flowing silver hair, adorned with luminescent crystals, intricate crown of twisted gold and moonstone, soft ethereal lighting, detailed facial features, fantasy art style, highly detailed, painted by Artgerm and Charlie Bowater"</p>
-                    </div>
-                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
-                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌃 Cyberpunk Scene</p>
-                        <p style="margin: 0; font-style: italic;">"Neon-lit cyberpunk street market in rain, holographic advertisements reflecting in puddles, street vendors with glowing cyber-augmentations, dense urban environment, atmospheric fog, cinematic lighting, inspired by Blade Runner 2049"</p>
-                    </div>
-                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
-                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🎨 Abstract Art</p>
-                        <p style="margin: 0; font-style: italic;">"Vibrant abstract composition of flowing liquid colors, dynamic swirls of iridescent purples and teals, golden geometric patterns emerging from chaos, luxury art style, ultra-detailed, painted in oil on canvas, inspired by James Jean and Gustav Klimt"</p>
-                    </div>
-                    <div style="background: rgba(75, 121, 161, 0.1); padding: 1em; border-radius: 8px; margin-bottom: 1em;">
-                        <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌿 Macro Nature</p>
-                        <p style="margin: 0; font-style: italic;">"Extreme macro photography of a dewdrop on a butterfly wing, rainbow light refraction, crystalline clarity, intricate wing scales visible, natural bokeh background, professional studio lighting, shot with Canon MP-E 65mm lens"</p>
-                    </div>
-                </div>
-            """)
+                    
+                    # Examples section
+                    gr.Markdown('<p class="neon-text">Click on any example to use it:</p>')
+                    with gr.Row():
+                        example_boxes = [gr.Button(ex[:40] + "...", elem_classes=["neon-example"]) for ex in example_prompts]
+                    
+                    # Connect example buttons to the prompt textbox
+                    for i, example_btn in enumerate(example_boxes):
+                        example_btn.click(
+                            fn=lambda x=example_prompts[i]: x,
+                            outputs=prompt
+                        )
+                    
+                    with gr.Accordion("Advanced Settings", open=False):
+                        with gr.Group():
+                            with gr.Row():
+                                height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024, 
+                                                elem_classes=["neon-slider"])
+                                width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024,
+                                                elem_classes=["neon-slider"])
+                            
+                            with gr.Row():
+                                steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8,
+                                                elem_classes=["neon-slider"])
+                                scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5,
+                                                elem_classes=["neon-slider"])
+                            
+                            seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0,
+                                            elem_classes=["neon-input"])
+                    
+                    generate_btn = gr.Button("Generate Image", variant="primary", scale=1, elem_classes=["neon-button"])
+
+            with gr.Column(scale=4, elem_classes=["neon-card"]):
+                output = gr.Image(label="Your Generated Image")
+        
+        gr.Markdown(
+            """
+            <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px;" class="neon-card">
+                <h2 style="font-size: 1.5rem; margin-bottom: 1rem;" class="neon-text">How to Use</h2>
+                <ol style="padding-left: 1.5rem; color: #00f3ff;">
+                    <li>Enter a detailed description of the image you want to create.</li>
+                    <li>Or click one of our exciting example prompts above!</li>
+                    <li>Adjust advanced settings if desired (tap to expand).</li>
+                    <li>Tap "Generate Image" and wait for your creation!</li>
+                </ol>
+                <p style="margin-top: 1rem; font-style: italic; color: #ff00dd;">Tip: Be specific in your description for best results!</p>
+            </div>
+            """
+        )
 
-        with gr.Column(scale=4, elem_classes=["fixed-width"]):
-            output = gr.Image(
-                label="Generated Image",
-                elem_id="output-image",
-                elem_classes=["output-image", "fixed-width"]
-            )
-    
     @spaces.GPU
     def process_image(height, width, steps, scales, prompt, seed):
         global pipe
-        
-        # 입력값 검증
-        if not prompt or prompt.strip() == "":
-            return None, "", False, "이미지 설명을 입력해주세요.", True
-            
-        # 프롬프트 필터링
-        is_safe, filtered_prompt = filter_prompt(prompt)
-        if not is_safe:
-            return None, "", False, "부적절한 내용이 포함된 프롬프트입니다.", True
-            
-        try:
-            # 메모리 확보를 위한 가비지 콜렉션
-            gc.collect()
-            torch.cuda.empty_cache()
-            
-            # 시드 값 확인 및 보정
-            if seed is None or not isinstance(seed, (int, float)):
-                seed = get_random_seed()
-            else:
-                seed = int(seed)  # 타입 변환 안전하게 처리
-                
-            # 높이와 너비를 64의 배수로 조정 (FLUX 모델 요구사항)
-            height = (int(height) // 64) * 64
-            width = (int(width) // 64) * 64
-            
-            # 안전장치 - 최대값 제한
-            steps = min(int(steps), 25)
-            scales = max(min(float(scales), 5.0), 0.0)
-            
-            # 이미지 생성
-            with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-                # 중요: generator 설정 시 device를 명시적으로 지정
-                generator = torch.Generator(device="cuda").manual_seed(seed)
-                
-                # 모든 텐서가 같은 디바이스에 있는지 확인
-                for name, module in pipe.components.items():
-                    if hasattr(module, 'device') and module.device.type != "cuda":
-                        module.to("cuda")
-                        
-                # 이미지 생성 - 모든 매개변수에 device를 명시적 지정
-                generated_image = pipe(
-                    prompt=[filtered_prompt],
-                    generator=generator,
-                    num_inference_steps=steps,
-                    guidance_scale=scales,
-                    height=height,
-                    width=width,
-                    max_sequence_length=256,
-                    device="cuda"  # 명시적 device 지정
-                ).images[0]
-                
-                return generated_image, "", False, "", False
-                
-        except Exception as e:
-            error_msg = f"이미지 생성 중 오류가 발생했습니다: {str(e)}"
-            print(error_msg)
-            traceback.print_exc()
-            
-            # 오류 후 메모리 정리
-            gc.collect()
-            torch.cuda.empty_cache()
-            
-            return None, "", False, error_msg, True
-    
-    def update_seed():
-        return get_random_seed()
-    
-    # 이미지 생성 준비 함수
-    def prepare_generation(height, width, steps, scales, prompt, seed):
-        global pipe
-        
-        # 모델이 아직 로드되지 않았다면 로드
-        if pipe is None:
-            # 로딩 상태 표시
-            loading_message = "모델을 로딩 중입니다... 처음 실행 시 시간이 소요될 수 있습니다."
-            
-            is_loaded = initialize_model()
-            if not is_loaded:
-                return None, "", False, "모델 로딩에 실패했습니다. 페이지를 새로고침하고 다시 시도해 주세요.", True
-            
-        # 생성 상태 표시
-        loading_message = "이미지를 생성 중입니다..."
-            
-        # 생성 프로세스 시작
-        return process_image(height, width, steps, scales, prompt, seed)
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+            return pipe(
+                prompt=[prompt],
+                generator=torch.Generator().manual_seed(int(seed)),
+                num_inference_steps=int(steps),
+                guidance_scale=float(scales),
+                height=int(height),
+                width=int(width),
+                max_sequence_length=256
+            ).images[0]
 
-    # 버튼 클릭 이벤트 연결
     generate_btn.click(
-        fn=prepare_generation,
+        process_image,
         inputs=[height, width, steps, scales, prompt, seed],
-        outputs=[output, loading_status, loading_status, error_message, error_message]
-    )
-    
-    randomize_seed.click(
-        fn=update_seed,
-        outputs=[seed]
+        outputs=output
     )
 
 if __name__ == "__main__":
-
-    demo.queue(max_size=10).launch()
+    demo.launch()