Update app.py

app.py

@@ -5,7 +5,7 @@ import os
 from compel import Compel, ReturnedEmbeddingsType
 from diffusers import DiffusionPipeline
 
-
+# 모델 로드
 model_name = os.environ.get('MODEL_NAME', 'UnfilteredAI/NSFW-gen-v2')
 pipe = DiffusionPipeline.from_pretrained(
     model_name,
@@ -13,19 +13,34 @@ pipe = DiffusionPipeline.from_pretrained(
 )
 pipe.to('cuda')
 
+# Compel 설정
 compel = Compel(
-  tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,
+  tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
   text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
   returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
   requires_pooled=[False, True]
 )
 
+# 기본 네거티브 프롬프트
+default_negative_prompt = "(low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn, (deformed | distorted | disfigured:1.3), bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers:1.4, disconnected limbs, blurry, amputation."
+
+# 예제 프롬프트 목록
+example_prompts = [
+    ["a beautiful woman in a summer dress at the beach, golden sunset, professional photography, 8k", default_negative_prompt, 40, 7.5, 1024, 1024, 4],
+    ["portrait of a cyberpunk character, highly detailed, neon lights, futuristic cityscape in background, 8k, ultra realistic", default_negative_prompt, 50, 8.0, 768, 1024, 4],
+    ["detailed fantasy art of magical forest with fairies, ethereal lighting, mystical atmosphere", default_negative_prompt, 60, 7.0, 1024, 768, 4],
+    ["photorealistic portrait of a stunning model, studio lighting, fashion photography", default_negative_prompt, 45, 7.0, 1024, 1024, 4],
+]
 
+# 이미지 생성 함수
 @spaces.GPU(duration=120)
-def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples):
+def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples, progress=gr.Progress()):
+    progress(0, desc="Preparing")
     embeds, pooled = compel(prompt)
     neg_embeds, neg_pooled = compel(negative_prompt)
-    return pipe(
+    
+    progress(0.1, desc="Generating images")
+    images = pipe(
         prompt_embeds=embeds,
         pooled_prompt_embeds=pooled,
         negative_prompt_embeds=neg_embeds,
@@ -34,20 +49,116 @@ def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width
         guidance_scale=guidance_scale,
         width=width,
         height=height,
-        num_images_per_prompt=num_samples
+        num_images_per_prompt=num_samples,
+        callback=lambda i, t, latents: progress((i + 1) / num_inference_steps)
     ).images
+    
+    return images
+
+# CSS 스타일
+css = """
+.gallery-item {
+    transition: transform 0.2s;
+    box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+    border-radius: 10px;
+}
+.gallery-item:hover {
+    transform: scale(1.03);
+    box-shadow: 0 8px 16px rgba(0,0,0,0.2);
+}
+.container {
+    max-width: 1200px;
+    margin: auto;
+}
+.header {
+    text-align: center;
+    margin-bottom: 2rem;
+    padding: 1rem;
+    background: linear-gradient(90deg, rgba(76,0,161,0.8) 0%, rgba(28,110,164,0.8) 100%);
+    border-radius: 10px;
+    color: white;
+}
+.slider-container {
+    background-color: #f5f5f5;
+    padding: 1rem;
+    border-radius: 10px;
+    margin-bottom: 1rem;
+}
+.prompt-container {
+    background-color: #f0f8ff;
+    padding: 1rem;
+    border-radius: 10px;
+    margin-bottom: 1rem;
+    border: 1px solid #d0e8ff;
+}
+.examples-header {
+    background: linear-gradient(90deg, rgba(41,128,185,0.7) 0%, rgba(142,68,173,0.7) 100%);
+    color: white;
+    padding: 0.5rem;
+    border-radius: 8px;
+    text-align: center;
+    margin-bottom: 0.5rem;
+}
+"""
 
+# Gradio 인터페이스
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    gr.HTML("""
+    <div class="header">
+        <h1>🎨 UnfilteredAI NSFW-gen-v2 이미지 생성기</h1>
+        <p>창의적인 프롬프트를 입력하고 고품질 이미지를 생성해보세요.</p>
+    </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Group(elem_classes="prompt-container"):
+                prompt = gr.Textbox(label="프롬프트", placeholder="원하는 이미지를 설명해주세요...", lines=3)
+                negative_prompt = gr.Textbox(
+                    label="네거티브 프롬프트", 
+                    value=default_negative_prompt,
+                    lines=3
+                )
+            
+            with gr.Group(elem_classes="slider-container"):
+                with gr.Row():
+                    with gr.Column():
+                        steps = gr.Slider(minimum=20, maximum=100, value=60, step=1, label="추론 단계 (Quality)", info="높을수록 품질 향상 (시간 증가)")
+                        guidance = gr.Slider(minimum=1, maximum=15, value=7, step=0.1, label="가이던스 스케일 (Creativity)", info="낮을수록 더 창의적인 결과")
+                    
+                    with gr.Column():
+                        with gr.Row():
+                            width = gr.Slider(minimum=512, maximum=1536, value=1024, step=128, label="너비")
+                            height = gr.Slider(minimum=512, maximum=1536, value=1024, step=128, label="높이")
+                        
+                        num_samples = gr.Slider(minimum=1, maximum=8, value=4, step=1, label="생성할 이미지 수", info="한 번에 생성할 이미지 수")
+            
+            generate_btn = gr.Button("🚀 이미지 생성", variant="primary", size="lg")
+        
+        with gr.Column(scale=3):
+            output_gallery = gr.Gallery(label="생성된 이미지", elem_classes="gallery-item", columns=2, object_fit="contain", height=650)
+    
+    gr.HTML("""<div class="examples-header"><h3>✨ 예제 프롬프트</h3></div>""")
+    gr.Examples(
+        examples=example_prompts,
+        inputs=[prompt, negative_prompt, steps, guidance, width, height, num_samples],
+        outputs=output_gallery,
+        fn=generate,
+        cache_examples=True,
+    )
+    
+    # 이벤트 연결
+    generate_btn.click(
+        fn=generate,
+        inputs=[prompt, negative_prompt, steps, guidance, width, height, num_samples],
+        outputs=output_gallery
+    )
+    
+    gr.HTML("""
+    <div style="text-align: center; margin-top: 20px; padding: 10px; background-color: #f0f0f0; border-radius: 10px;">
+        <p>💡 팁: 고품질 이미지를 위해서는 세부적인 프롬프트와 더 높은 추론 단계를 사용하세요.</p>
+        <p>예) "professional photography, 8k, highly detailed, sharp focus, HDR" 등의 품질 관련 용어를 추가해보세요.</p>
+    </div>
+    """)
 
-gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Text(label="Prompt"),
-        gr.Text("(low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn, (deformed | distorted | disfigured:1.3), bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers:1.4, disconnected limbs, blurry, amputation.", label="Negative Prompt"),
-        gr.Number(60, label="Number inference steps"),
-        gr.Number(7, label="Guidance scale"),
-        gr.Number(1024, label="Width"),
-        gr.Number(1024, label="Height"),
-        gr.Number(7, label="# images"),
-    ],
-    outputs=gr.Gallery(),
-).launch()
+demo.launch()