Update app.py

app.py

@@ -1,131 +1,112 @@
 import os
+import torch
+import random
+import importlib
+from PIL import Image
 from huggingface_hub import snapshot_download
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForCausalLM, CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
+from diffusers import StableDiffusionPipeline, DiffusionPipeline, EulerDiscreteScheduler, UNet2DConditionModel
 
-# 모델 리비전과 로컬 저장 경로 설정
+# 환경 설정
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 REVISION = "ceaf371f01ef66192264811b390bccad475a4f02"
-LOCAL_FLORENCE = snapshot_download(
-    repo_id="microsoft/Florence-2-base",
-    revision=REVISION
-)
-LOCAL_TURBOX = snapshot_download(
-    repo_id="tensorart/stable-diffusion-3.5-large-TurboX"
-)
 
-import sys, types, importlib.machinery, importlib
-# flash_attn 무효화 처리
-spec = importlib.machinery.ModuleSpec('flash_attn', loader=None)
-mod = types.ModuleType('flash_attn')
-mod.__spec__ = spec
-sys.modules['flash_attn'] = mod
+# 로컬 다운로드
+LOCAL_FLORENCE = snapshot_download("microsoft/Florence-2-base", revision=REVISION)
+LOCAL_TURBOX = snapshot_download("tensorart/stable-diffusion-3.5-large-TurboX")
 
-import gradio as gr
-import torch
-import random
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
-from transformers import (    CLIPTextModel,    CLIPTokenizer,    CLIPFeatureExtractor)
-from diffusers import DiffusionPipeline, EulerDiscreteScheduler
+# 디바이스 및 dtype 설정
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# 1. 경량화 옵션: FP16 + 8bit 양자화 적용
-model_repo = "tensorart/stable-diffusion-3.5-large-TurboX"
-pipe = DiffusionPipeline.from_pretrained(
-    model_repo,
-    torch_dtype=torch.float16,
-    load_in_8bit=True,
-    device_map="auto",
-    safety_checker=None,
-    feature_extractor=None
-)
-pipe = pipe.to("cuda")  # GPU 사용
-# 스케줄러 로드 (Euler 방식)
-pipe.scheduler = EulerDiscreteScheduler.from_pretrained(
-    model_repo, subfolder="scheduler", local_files_only=True
+# 모델 로딩 (부분별 로딩 + dtype 적용)
+scheduler = EulerDiscreteScheduler.from_pretrained(
+    LOCAL_TURBOX, subfolder="scheduler", torch_dtype=dtype
 )
+text_encoder = CLIPTextModel.from_pretrained(LOCAL_TURBOX, subfolder="text_encoder", torch_dtype=dtype)
+tokenizer = CLIPTokenizer.from_pretrained(LOCAL_TURBOX, subfolder="tokenizer")
+feature_extractor = CLIPFeatureExtractor.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="feature_extractor")
+unet = UNet2DConditionModel.from_pretrained(LOCAL_TURBOX, subfolder="unet", torch_dtype=dtype)
 
-# Florence 모델 설정 (CPU 로드 후 필요 시 GPU 이동)
 florence_model = AutoModelForCausalLM.from_pretrained(
-    LOCAL_FLORENCE,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-    load_in_8bit=True  # 텍스트 생성 모델도 8bit 로드
+    LOCAL_FLORENCE, trust_remote_code=True, torch_dtype=dtype
 )
-florence_model = florence_model.to("cpu")
-florence_model.eval()
-florence_processor = AutoProcessor.from_pretrained(
-    LOCAL_FLORENCE,
-    trust_remote_code=True
+florence_model.to("cpu").eval()
+florence_processor = AutoProcessor.from_pretrained(LOCAL_FLORENCE, trust_remote_code=True)
+
+# Stable Diffusion 파이프라인
+pipe = DiffusionPipeline.from_pretrained(
+    LOCAL_TURBOX,
+    torch_dtype=dtype,
+    trust_remote_code=True,
+    safety_checker=None,
+    feature_extractor=None
 )
+pipe = pipe.to(device)
+pipe.scheduler = scheduler
+pipe.enable_attention_slicing()  # 메모리 절약
 
-# 최대 시드 값
+# 상수
 MAX_SEED = 2**31 - 1
 
-def preprocess_image(input_image: Image.Image) -> Image.Image:
-    """
-    입력 이미지 크기를 512x512로 리사이징
-    """
-    img = input_image.convert("RGB")
-    img = img.resize((512, 512), resample=Image.LANCZOS)  # 고품질 리사이징
-    return img
-
-# 영어 설명을 카툰 스타일 한국어 프롬프트로 변환
+# 텍스트 스타일러
 def pseudo_translate_to_korean_style(en_prompt: str) -> str:
-    return f"카툰 스타일: {en_prompt} 아름다운 인물"
+    return f"Cartoon styled {en_prompt} handsome or pretty people"
 
-# 이미지 → 상세 설명 → 카툰 프롬프트 생성
+# 프롬프트 생성
 def generate_prompt(image):
-    img = preprocess_image(image)
-    inputs = florence_processor(
-        text="<MORE_DETAILED_CAPTION>",
-        images=img,
-        return_tensors="pt"
-    ).to(pipe.device)
-    generated_ids = florence_model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=256,  # 토큰 수 제한
-        num_beams=2         # 빔 서치 개수 축소
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to("cpu")
+    with torch.no_grad():
+        generated_ids = florence_model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=256,
+            num_beams=3
+        )
+    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = florence_processor.post_process_generation(
+        generated_text,
+        task="<MORE_DETAILED_CAPTION>",
+        image_size=(image.width, image.height)
     )
-    generated_text = florence_processor.batch_decode(
-        generated_ids, skip_special_tokens=True
-    )[0]
-    return pseudo_translate_to_korean_style(generated_text)
+    prompt_en = parsed_answer["<MORE_DETAILED_CAPTION>"]
+    cartoon_prompt = pseudo_translate_to_korean_style(prompt_en)
+    return cartoon_prompt
 
-# 프롬프트 → 이미지 생성
+# 이미지 생성 함수
 def generate_image(prompt, seed=42, randomize_seed=False):
-    # 시드 랜덤화 옵션
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device="cuda").manual_seed(seed)
-    # 해상도를 512x512로 설정해 메모리 절약
+    generator = torch.Generator().manual_seed(seed)
     image = pipe(
         prompt=prompt,
-        guidance_scale=1.2,
-        num_inference_steps=8,
+        guidance_scale=1.5,
+        num_inference_steps=6,  # 최적화된 step 수
         width=512,
         height=512,
         generator=generator
     ).images[0]
     return image, seed
 
-# Gradio UI 구성
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🖼 이미지 → 설명 → 카툰 이미지 생성기")
-    gr.Markdown(
-        "**사용법 (한글 설명)**\n"
-        "1. 왼쪽에 이미지를 업로드하세요.\n"
-        "2. AI가 상세 설명을 생성하고 카툰 스타일 한국어 프롬프트로 변환합니다.\n"
-        "3. 오른쪽에 생성된 카툰 이미지를 확인하세요."
-    )
+    gr.Markdown("# 🖼 이미지 → 설명 생성 → 카툰 이미지 자동 생성기")
+    gr.Markdown("**📌 사용법 안내 (한국어)**\n"
+                "- 이미지를 업로드하면 AI가 설명 → 스타일 변환 → 카툰 이미지 생성까지 자동으로 수행합니다.")
+
     with gr.Row():
         with gr.Column():
-            input_img = gr.Image(type="pil", label="원본 이미지 업로드")  # 타입: PIL 이미지
-            run_button = gr.Button("생성 시작")
+            input_img = gr.Image(label="🎨 원본 이미지 업로드")
+            run_button = gr.Button("✨ 생성 시작")
+
         with gr.Column():
-            prompt_out = gr.Textbox(label="생성된 프롬프트", lines=2, show_copy_button=True)
-            output_img = gr.Image(label="생성된 카툰 이미지")
+            prompt_out = gr.Textbox(label="📝 스타일 적용된 프롬프트", lines=3, show_copy_button=True)
+            output_img = gr.Image(label="🎉 생성된 이미지")
 
-    # 버튼 클릭 시 전체 프로세스 실행
     def full_process(img):
         prompt = generate_prompt(img)
         image, seed = generate_image(prompt, randomize_seed=True)
@@ -133,4 +114,4 @@ with gr.Blocks() as demo:
 
     run_button.click(fn=full_process, inputs=[input_img], outputs=[prompt_out, output_img])
 
-demo.launch(share=True)  # Hugging Face Spaces에 배포 시 share=True 사용
+demo.launch()