Update app.py

app.py

@@ -24,7 +24,9 @@ import gradio as gr
 import numpy as np
 import requests
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline
+from diffusers import EulerAncestralDiscreteScheduler
+from compel import Compel, ReturnedEmbeddingsType
 from PIL import Image
 
 # ===== OpenAI 설정 =====
@@ -52,24 +54,45 @@ if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
 # ===== 디바이스 & 모델 로드 =====
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-repo_id = "black-forest-labs/FLUX.1-dev"
-adapter_id = "seawolf2357/kim-korea"
-
 # Add error handling for model loading
 try:
-    pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-    pipeline.load_lora_weights(adapter_id)
-    pipeline = pipeline.to(device)
+    # Make sure to use torch.float16 consistently throughout the pipeline
+    pipeline = StableDiffusionXLPipeline.from_pretrained(
+        "votepurchase/pornmasterPro_noobV3VAE",
+        torch_dtype=torch.float16,
+        variant="fp16",  # Explicitly use fp16 variant
+        use_safetensors=True  # Use safetensors if available
+    )
+    
+    pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
+    pipeline.to(device)
+    
+    # Force all components to use the same dtype
+    pipeline.text_encoder.to(torch.float16)
+    pipeline.text_encoder_2.to(torch.float16)
+    pipeline.vae.to(torch.float16)
+    pipeline.unet.to(torch.float16)
+    
+    # Initialize Compel for long prompt processing
+    compel = Compel(
+        tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
+        text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
+        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        requires_pooled=[False, True],
+        truncate_long_prompts=False
+    )
+    
     print("Model loaded successfully")
 except Exception as e:
     print(f"Error loading model: {e}")
     pipeline = None
+    compel = None
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 1216
 
 # ===== 한글 여부 판별 =====
 HANGUL_RE = re.compile(r"[\u3131-\u318E\uAC00-\uD7A3]+")
@@ -182,23 +205,64 @@ def save_generated_image(image: Image.Image, prompt: str) -> str:
         f.write(f"{filename}|{prompt}|{timestamp}\n")
     return filepath
 
+# ===== Long prompt processing function =====
+def process_long_prompt(prompt, negative_prompt=""):
+    """Simple long prompt processing using Compel"""
+    try:
+        conditioning, pooled = compel([prompt, negative_prompt])
+        return conditioning, pooled
+    except Exception as e:
+        print(f"Long prompt processing failed: {e}, falling back to standard processing")
+        return None, None
+
 # ===== Diffusion 호출 =====
 
-def run_pipeline(prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int, lora_scale: float):
+def run_pipeline(prompt: str, negative_prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int):
     if pipeline is None:
         raise ValueError("Model pipeline not loaded")
     
     generator = torch.Generator(device=device).manual_seed(int(seed))
-    result = pipeline(
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_steps,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-    ).images[0]
-    return result
+    
+    # Check if prompt is long
+    use_long_prompt = len(prompt.split()) > 60 or len(prompt) > 300
+    
+    try:
+        # Try long prompt processing first if prompt is long
+        if use_long_prompt and compel is not None:
+            print("Using long prompt processing...")
+            conditioning, pooled = process_long_prompt(prompt, negative_prompt)
+            
+            if conditioning is not None:
+                result = pipeline(
+                    prompt_embeds=conditioning[0:1],
+                    pooled_prompt_embeds=pooled[0:1],
+                    negative_prompt_embeds=conditioning[1:2],
+                    negative_pooled_prompt_embeds=pooled[1:2],
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_steps,
+                    width=width,
+                    height=height,
+                    generator=generator
+                ).images[0]
+                return result
+        
+        # Fall back to standard processing
+        result = pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_steps,
+            width=width,
+            height=height,
+            generator=generator
+        ).images[0]
+        
+        return result
+    except RuntimeError as e:
+        print(f"Error during generation: {e}")
+        # Return a blank image with error message
+        error_img = Image.new('RGB', (width, height), color=(0, 0, 0))
+        return error_img
 
 # ===== Gradio inference 래퍼 =====
 
@@ -207,13 +271,13 @@ def generate_image(
     user_prompt: str,
     style_key: str,
     enhance_prompt_enabled: bool = False,
+    negative_prompt: str = "nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn",
     seed: int = 42,
     randomize_seed: bool = True,
     width: int = 1024,
-    height: int = 768,
-    guidance_scale: float = 3.5,
-    num_inference_steps: int = 30,
-    lora_scale: float = 1.0,
+    height: int = 1024,
+    guidance_scale: float = 7.0,
+    num_inference_steps: int = 28,
     progress=None,
 ):
     try:
@@ -225,7 +289,7 @@ def generate_image(
         print(f"Final prompt: {final_prompt}")
 
         # 2) 파이프라인 호출
-        image = run_pipeline(final_prompt, seed, width, height, guidance_scale, num_inference_steps, lora_scale)
+        image = run_pipeline(final_prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps)
 
         # 3) 저장
         save_generated_image(image, final_prompt)
@@ -241,17 +305,16 @@ def generate_image(
 # ===== 예시 프롬프트 (한국어/영어 혼용 허용) =====
 
 examples = [
-    "Mr. KIM이 두 손으로 'Fighting!' 현수막을 들고 있는 모습, 애국심과 국가 발전에 대한 의지를 보여주고 있다.",   
-    "Mr. KIM이 양팔을 들어 올리며 승리의 표정으로 환호하는 모습, 승리와 미래에 대한 희망을 보여주고 있다.",
-    "Mr. KIM이 운동복을 입고 공원에서 조깅하는 모습, 건강한 생활습관과 활기찬 리더십을 보여주고 있다.",  
-    "Mr. KIM이 붐비는 거리에서 여성 시민들과 따뜻하게 악수하는 모습, 여성 유권자들에 대한 진정한 관심과 소통을 보여주고 있다.",
-    "Mr. KIM이 선거 유세장에서 지평선을 향해 손가락으로 가리키며 영감을 주는 제스처를 취하고 있고, 여성들과 아이들이 박수를 치고 있다.",
-    "Mr. KIM이 지역 행사에 참여하여 열정적으로 응원하는 여성 지지자들에게 둘러싸여 있는 모습.",
-    "Mr. KIM visiting a local market, engaging in friendly conversation with female vendors and shopkeepers.",
-    "Mr. KIM walking through a university campus, discussing education policies with female students and professors.",    
-    "Mr. KIM delivering a powerful speech in front of a large crowd with confident gestures and determined expression.",
-    "Mr. KIM in a dynamic interview setting, passionately outlining his visions for the future.",
-    "Mr. KIM preparing for an important debate, surrounded by paperwork, looking focused and resolute.",
+    "아름다운 한국 여성이 전통 한복을 입고 경복궁 앞에서 미소 짓는 모습",   
+    "현대적인 서울 도시 풍경과 남산타워가 보이는 야경",
+    "한국의 봄, 벚꽃이 만개한 공원에서 산책하는 사람들",  
+    "전통 한옥마을의 고즈넉한 오후 풍경",
+    "K-pop 아이돌 콘서트의 화려한 무대와 열정적인 팬들",
+    "서울의 번화한 명동 거리, 쇼핑하는 사람들과 네온사인",
+    "Beautiful Korean woman in traditional hanbok dress standing in front of Gyeongbokgung Palace",
+    "Modern Seoul cityscape at night with N Seoul Tower illuminated",    
+    "Korean spring scenery with cherry blossoms in full bloom",
+    "Traditional Korean tea ceremony in a peaceful hanok setting",
 ]
 
 # ===== 커스텀 CSS (진한 붉은색 고급 디자인) =====
@@ -571,8 +634,8 @@ def create_interface():
         with gr.Group(elem_classes="model-description"):
             gr.HTML("""
             <p>
-            <strong>Mr. KIM in KOREA</strong><br>
-            <small style="opacity: 0.8;">본 모델은 연구 목적으로 특정인의 얼굴과 외모를 LoRA 기술로 학습한 모델입니다.목적 외의 용도로 무단 사용하지 않도록 유의해 주세요. 프롬프트에 'kim'을 포함하여 주세요.</small><br><br>
+            <strong>고품질 이미지 생성 AI</strong><br>
+            <small style="opacity: 0.8;">StableDiffusionXL 기반의 고급 이미지 생성 모델입니다. 한국어 프롬프트를 자동으로 번역하며, 긴 프롬프트도 지원합니다.</small><br><br>
             """)
 
         # ===== 메인 입력 =====
@@ -615,15 +678,21 @@ def create_interface():
 
             # ===== 고급 설정 =====
             with gr.Accordion("Advanced Settings (고급 설정)", open=False, elem_classes="advanced-settings"):
+                negative_prompt = gr.Text(
+                    label="Negative prompt (네거티브 프롬프트)",
+                    max_lines=1,
+                    placeholder="Enter a negative prompt",
+                    value="nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"
+                )
+                
                 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)
+                    height = gr.Slider(label="Height (세로)", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                 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="Inference steps (추론 단계)", minimum=1, maximum=50, step=1, value=30)
-                    lora_scale = gr.Slider(label="LoRA scale (LoRA 스케일)", minimum=0.0, maximum=1.0, step=0.1, value=1.0)
+                    guidance_scale = gr.Slider(label="Guidance scale (가이던스 스케일)", minimum=0.0, maximum=20.0, step=0.1, value=7.0)
+                    num_inference_steps = gr.Slider(label="Inference steps (추론 단계)", minimum=1, maximum=50, step=1, value=28)
 
             # ===== 예시 영역 =====
             with gr.Group(elem_classes="example-region"):
@@ -637,13 +706,13 @@ def create_interface():
                 user_prompt,
                 style_select,
                 enhance_prompt_checkbox,
+                negative_prompt,
                 seed,
                 randomize_seed,
                 width,
                 height,
                 guidance_scale,
                 num_inference_steps,
-                lora_scale,
             ],
             outputs=[result_image, seed_output],
         )