add : refinement logic

app.py

@@ -7,12 +7,19 @@ from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, Autoe
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
+from llm_wrapper import run_gemini
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 import subprocess
 
 subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
+
+from pydantic import BaseModel
+
+class RefinedPrompt(BaseModel):
+    prompt: str
+
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -38,9 +45,15 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
+
+    refined_prompt = run_gemini(
+        target_prompt=prompt,
+        prompt_in_path="prompt.json",
+        output_structure=RefinedPrompt,
+    )
     
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt,
+            prompt=refined_prompt.prompt,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             width=width,
@@ -52,9 +65,9 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
             yield img, seed
     
 examples = [
-    "로켓에 타고 있는 포닉스",
-    "일렉기타를 들고 있는 포닉스",
-    "컴퓨터공학을 공부중인 포닉스",
+    "기계공학과(로켓) 포닉스",
+    "바이올린을 연주하는 포닉스",
+    "물리학을 연구하는 포닉스",
 ]
 
 css="""
@@ -67,7 +80,7 @@ css="""
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# [POSTECH] PONIX Generator
+        gr.Markdown(f"""# [POSTECH] PONIX Generator 🌊
 [[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
         """)
         

llm_wrapper.py

@@ -0,0 +1,107 @@
+import logging
+from PIL import Image
+from io import BytesIO
+import requests, os, json, time
+
+from google import genai
+
+prompt_base_path = "src/llm_wrapper/prompt"
+client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+
+
+def encode_image(image_source):
+    """
+    이미지 경로가 URL이든 로컬 파일이든 Pillow Image 객체이든 동일하게 처리하는 함수.
+    이미지를 열어 google.genai.types.Part 객체로 변환합니다.
+    Pillow에서 지원되지 않는 포맷에 대해서는 예외를 발생시킵니다.
+    """
+    try:
+        # 이미 Pillow 이미지 객체인 경우 그대로 사용
+        if isinstance(image_source, Image.Image):
+            image = image_source
+        else:
+            # URL에서 이미지 다운로드
+            if isinstance(image_source, str) and (
+                image_source.startswith("http://")
+                or image_source.startswith("https://")
+            ):
+                response = requests.get(image_source)
+                image = Image.open(BytesIO(response.content))
+            # 로컬 파일에서 이미지 열기
+            else:
+                image = Image.open(image_source)
+
+        # 이미지 포맷이 None인 경우 (메모리에서 생성된 이미지 등)
+        if image.format is None:
+            image_format = "JPEG"
+        else:
+            image_format = image.format
+
+        # 이미지 포맷이 지원되지 않는 경우 예외 발생
+        if image_format not in Image.registered_extensions().values():
+            raise ValueError(f"Unsupported image format: {image_format}.")
+
+        buffered = BytesIO()
+        # PIL에서 지원되지 않는 포맷이나 다양한 채널을 RGB로 변환 후 저장
+        if image.mode in ("RGBA", "P", "CMYK"):  # RGBA, 팔레트, CMYK 등은 RGB로 변환
+            image = image.convert("RGB")
+        image.save(buffered, format="JPEG")
+        
+        return genai.types.Part.from_bytes(data=buffered.getvalue(), mime_type="image/jpeg")
+
+    except requests.exceptions.RequestException as e:
+        raise ValueError(f"Failed to download the image from URL: {e}")
+    except IOError as e:
+        raise ValueError(f"Failed to process the image file: {e}")
+    except ValueError as e:
+        raise ValueError(e)
+
+
+def run_gemini(
+    target_prompt: str,
+    prompt_in_path: str,
+    output_structure,
+    img_in_data: str = None,
+    model: str = "gemini-2.0-flash",
+) -> str:
+    """
+    GEMINI API를 동기 방식으로 호출하여 문자열 응답을 받습니다.
+    retry 논리는 제거되었습니다.
+    """
+    with open(os.path.join(prompt_base_path, prompt_in_path), "r", encoding="utf-8") as file:
+        prompt_dict = json.load(file)
+
+    system_prompt = prompt_dict["system_prompt"]
+    user_prompt_head = prompt_dict["user_prompt"]["head"]
+    user_prompt_tail = prompt_dict["user_prompt"]["tail"]
+
+    user_prompt_text = "\n".join([user_prompt_head, target_prompt, user_prompt_tail])
+    input_content = [user_prompt_text]
+
+    if img_in_data is not None:
+        encoded_image = encode_image(img_in_data)
+        input_content.append(encoded_image)
+        
+    logging.info("Requested API for chat completion response (sync call)...")
+    start_time = time.time()
+
+    # 동기 방식: client.models.generate_content(...)
+    chat_completion = client.models.generate_content(
+        model=model,
+        contents=input_content,
+        config={
+            "system_instruction": system_prompt,
+            "response_mime_type": "application/json",
+            "response_schema": output_structure,
+        }
+    )
+
+    chat_output = chat_completion.parsed
+    input_token = chat_completion.usage_metadata.prompt_token_count 
+    output_token = chat_completion.usage_metadata.candidates_token_count
+    pricing = input_token / 1000000 * 0.1 * 1500 + output_token / 1000000 * 0.7 * 1500
+
+    logging.info(
+        f"[GEMINI] Request completed (sync). Time taken: {time.time()-start_time:.2f}s / Pricing(KRW): {pricing:.2f}"
+    )
+    return chat_output, chat_completion

prompt.json

@@ -0,0 +1,7 @@
+{
+    "system_prompt": "**역할(Role)**\n\n당신(시스템)은 사용자가 입력한 요구사항(한국어 문장)을 받아,\nTextual Inversion이 적용된 토큰(\"<s0><s1><s2> plush bird\")이 반드시 포함된 영어 텍스트 프롬프트를 생성해야 합니다.\n추가로 배경·상황·스타일을 풍부하게 묘사하여, 실제 텍스트 투 이미지 모델에 넣기만 하면 원하는 장면이 생성되도록 도움을 줍니다.\n\n\n**주요 규칙(Rules)**\n\n**\"포닉스\"**라는 단어가 요구사항에 등장하면, 이를 **<s0><s1><s2> plush bird**로 치환한다.\n항상 프롬프트 문장 앞부분에 **photo of <s0><s1><s2> plush bird**를 포함한다.\n사용자가 원하는 **장면(배경, 환경, 상황)**을 영어로 자세히 묘사한다.\n스타일(e.g., hyper-realistic, cinematic lighting, 8k resolution, ultra high quality, 등)을 적절히 추가해 고퀄리티 이미지를 유도한다.\n출력은 단일 문자열(또는 여러 줄) 형태로 영어 문장 위주로 작성한다.\n\n\n\n**예시(Examples)**\n\nInput 1\n사용자: \"기계공학과(로켓) 포닉스\"\n시스템 변환:\nphoto of <s0><s1><s2> plush bird \nwearing an astronaut suit and space helmet\ninside a spacecraft cockpit during flight, \nsurrounded by control panels and navigation systems,\nblinking lights and monitoring screens,\nEarth visible through the spacecraft window in background,\nhyper-realistic details, cinematic lighting, 8k resolution, \nultra high quality photograph, \nhigh-tech space environment, adventurous atmosphere\n\n\nInput 2\n사용자: \"바이올린을 연주하는 포닉스\"\n시스템 변환:\nphoto of <s0><s1><s2> plush bird \nwearing an elegant black tailcoat,\ncrisp white dress shirt with bow tie,\nformal concert attire,\nholding a violin in playing position,\nin a grand concert hall with ornate architecture,\nwarm ambient lighting from chandeliers,\norchestra members visible in background,\nsheet music on stand nearby,\naudience in formal attire visible,\npolished wooden stage floor,\nconductor's podium visible in background,\nhyper-realistic details, warm classical lighting, 8k resolution,\nultra high quality photograph,\nprofessional classical concert environment, performance moment\n\n\nInput 3\n사용자: \"물리학을 연구하는 포닉스\"\n시스템 변환:\nphoto of <s0><s1><s2> plush bird \nwearing a lab coat and safety glasses,\ninside a physics laboratory,\nconducting experiments with quantum physics equipment,\nsurrounded by equations written on whiteboards,\noperating particle accelerator models,\nmeasuring devices and scientific instruments visible,\nhyper-realistic details, dramatic academic lighting, 8k resolution,\nultra high quality photograph,\nscientific environment, discovery atmosphere\n\n\n**출력 형식(Output Format)**\n\n최종 출력은 영어 텍스트로 된 하나의 프롬프트 문장(또는 여러 줄)이며,\n반드시 <s0><s1><s2> plush bird가 들어 있어야 함.\n상황에 따라 배경·디테일·조명·해상도를 다양한 형용사로 풍부히 기술해 준다.",
+    "user_prompt": {
+      "head": "",
+      "tail": ""
+    }
+}