Update app.py

app.py

@@ -1,83 +1,93 @@
-# 통합된 app.py 예시 (필수 라이브러리는 requirements.txt에 포함 필요)
-import gradio as gr
-import torch
-import random
-from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
-from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
-
-# 디바이스 설정
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Florence-2 로딩 (이미지 → 영어 설명 생성)
-florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
-
-# 번역기 (영어 → 한국어)
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ko", device=0 if torch.cuda.is_available() else -1)
-
-# SD 3.5 모델 로딩
-pipe = DiffusionPipeline.from_pretrained("tensorart/stable-diffusion-3.5-large-TurboX", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("tensorart/stable-diffusion-3.5-large-TurboX", subfolder="scheduler", shift=5)
-pipe = pipe.to(device)
-
-MAX_SEED = 2**31 - 1
-
-def generate_prompt(image):
-    """이미지에서 영어 캡션 추출 후 번역 및 스타일 변환"""
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-
-    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
-    generated_ids = florence_model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=512,
-        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)
-    )
-    prompt_en = parsed_answer["<MORE_DETAILED_CAPTION>"]
-    translated = translator(prompt_en, max_length=256)[0]['translation_text']
-    cartoon_prompt = f"{translated}, 카툰 스타일, 디지털 일러스트, 밝은 색감"
-    return cartoon_prompt
-
-def generate_image(prompt, seed=42, randomize_seed=False):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-        prompt=prompt,
-        negative_prompt="왜곡, 이상한 손, 비현실적 구조",
-        guidance_scale=1.5,
-        num_inference_steps=8,
-        width=768,
-        height=768,
-        generator=generator
-    ).images[0]
-    return image, seed
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 🖼 이미지 기반 자동 프롬프트 생성 및 한국어 카툰 스타일 변환")
-
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(label="왼쪽: 원본 이미지")
-            run_button = gr.Button("👉 프롬프트 생성 및 이미지 변환")
-
-        with gr.Column():
-            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)
-        return prompt, image
-
-    run_button.click(fn=full_process, inputs=[input_img], outputs=[prompt_out, output_img])
-
-demo.launch()
+import gradio as gr
+import torch
+import random
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
+from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+
+# Florence-2 로드 (이미지 캡셔닝)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to(device).eval()
+florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
+
+# 영어 → 한국어 번역기
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ko", device=0 if torch.cuda.is_available() else -1)
+
+# Stable Diffusion TurboX 로드
+model_repo = "tensorart/stable-diffusion-3.5-large-TurboX"
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo, subfolder="scheduler", shift=5)
+pipe = pipe.to(device)
+
+MAX_SEED = 2**31 - 1
+
+def generate_prompt(image):
+    """이미지 → 영어 설명 → 한국어 번역 + 카툰 스타일 적용"""
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
+    generated_ids = florence_model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=512,
+        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)
+    )
+    prompt_en = parsed_answer["<MORE_DETAILED_CAPTION>"]
+
+    # 영어 → 한국어 번역 후 카툰 스타일 추가
+    translated = translator(prompt_en, max_length=256)[0]['translation_text']
+    cartoon_prompt = f"{translated}, 카툰 스타일, 밝은 색감, 디지털 일러스트, 귀여운 분위기"
+    return cartoon_prompt
+
+def generate_image(prompt, seed=42, randomize_seed=False):
+    """텍스트 프롬프트 → 이미지 생성"""
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    image = pipe(
+        prompt=prompt,
+        negative_prompt="왜곡된 얼굴, 이상한 손, 흐림, 저해상도",
+        guidance_scale=1.5,
+        num_inference_steps=8,
+        width=768,
+        height=768,
+        generator=generator
+    ).images[0]
+    return image, seed
+
+# Gradio UI 정의
+with gr.Blocks() as demo:
+    gr.Markdown("# 🖼 이미지 → 한국어 프롬프트 → 카툰 스타일 이미지 생성기")
+    
+    gr.Markdown("**📌 사용법 안내 (한국어)**\n"
+                "- 왼쪽에 이미지를 업로드하세요.\n"
+                "- 자동으로 영어 프롬프트를 생성한 후, 한국어로 번역하고 카툰 스타일을 적용합니다.\n"
+                "- 오른쪽에 해당 프롬프트로 생성된 이미지가 출력됩니다.\n")
+
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="🎨 원본 이미지 업로드")
+            run_button = gr.Button("✨ 이미지 분석 및 생성 시작")
+
+        with gr.Column():
+            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)
+        return prompt, image
+
+    run_button.click(fn=full_process, inputs=[input_img], outputs=[prompt_out, output_img])
+
+demo.launch()