Update app.py

app.py

@@ -2,7 +2,7 @@ import spaces
 import random
 import torch
 from huggingface_hub import snapshot_download
-from transformers import CLIPProcessor, CLIPModel
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
 from kolors.pipelines import pipeline_stable_diffusion_xl_chatglm_256_ipadapter, pipeline_stable_diffusion_xl_chatglm_256
 from kolors.models.modeling_chatglm import ChatGLMModel
 from kolors.models.tokenization_chatglm import ChatGLMTokenizer
@@ -10,37 +10,34 @@ from kolors.models import unet_2d_condition
 from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 import gradio as gr
 import numpy as np
+from huggingface_hub import InferenceClient
 import os
-from accelerate import Accelerator
+from PIL import Image
+import re
 
-# Accelerate 초기화
-accelerator = Accelerator()
+# Cohere 모델 초기화
+client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))
 
-device = accelerator.device
-
-# 모델 체크포인트 다운로드
+device = "cuda"
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
 ckpt_IPA_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors-IP-Adapter-Plus")
 
-# ChatGLMModel, AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel 로드
 text_encoder = ChatGLMModel.from_pretrained(f'{ckpt_dir}/text_encoder', torch_dtype=torch.float16).half().to(device)
 tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
 vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half().to(device)
 scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")
 unet_t2i = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().to(device)
 unet_i2i = unet_2d_condition.UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet", revision=None).half().to(device)
-image_encoder = CLIPModel.from_pretrained(f'{ckpt_IPA_dir}/image_encoder', ignore_mismatched_sizes=True).to(dtype=torch.float16, device=device)
-
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(f'{ckpt_IPA_dir}/image_encoder', ignore_mismatched_sizes=True).to(dtype=torch.float16, device=device)
 ip_img_size = 336
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+clip_image_processor = CLIPImageProcessor(size=ip_img_size, crop_size=ip_img_size)
 
-# Stable Diffusion 파이프라인 설정
 pipe_t2i = pipeline_stable_diffusion_xl_chatglm_256.StableDiffusionXLPipeline(
     vae=vae,
-    text_encoder=text_encoder, 
-    tokenizer=tokenizer, 
-    unet=unet_t2i, 
-    scheduler=scheduler, 
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    unet=unet_t2i,
+    scheduler=scheduler,
     force_zeros_for_empty_prompt=False
 ).to(device)
 
@@ -51,7 +48,7 @@ pipe_i2i = pipeline_stable_diffusion_xl_chatglm_256_ipadapter.StableDiffusionXLP
     unet=unet_i2i,
     scheduler=scheduler,
     image_encoder=image_encoder,
-    feature_extractor=clip_processor,
+    feature_extractor=clip_image_processor,
     force_zeros_for_empty_prompt=False
 ).to(device)
 
@@ -60,23 +57,37 @@ if hasattr(pipe_i2i.unet, 'encoder_hid_proj'):
     
 pipe_i2i.load_ip_adapter(f'{ckpt_IPA_dir}' , subfolder="", weight_name=["ip_adapter_plus_general.bin"])
 
-# CLIP 모델 및 프로세서 로드
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
-clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def call_api(content, system_message, max_tokens=1000, temperature=0.7, top_p=0.95):
+    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
+    response = client.chat_completion(messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
+    return response.choices[0].message['content']
+
+def generate_prompt(korean_prompt):
+    system_message = """
+    Given the following description in Korean, 
+    translate and generate a concise English prompt suitable for a Stable Diffusion model. 
+    The prompt should be focused, descriptive, 
+    and contain specific keywords or phrases that will help guide the image generation process. 
+    Use simple and descriptive language, avoiding unnecessary words. 
+    Ensure the output is in English and follows the format typically used in Stable Diffusion prompts. 
+    The description is: [Insert Korean description here]
+    """
+    optimized_prompt = call_api(korean_prompt, system_message)
+    return optimized_prompt  # 최적화된 프롬프트 반환
+
 @spaces.GPU
-def infer(prompt, 
-          ip_adapter_image = None, 
-          ip_adapter_scale = 0.5, 
-          negative_prompt = "", 
-          seed = 0, 
-          randomize_seed = False, 
-          width = 1024, 
-          height = 1024, 
-          guidance_scale = 5.0, 
+def infer(prompt,
+          ip_adapter_image = None,
+          ip_adapter_scale = 0.5,
+          negative_prompt = "",
+          seed = 0,
+          randomize_seed = False,
+          width = 1024,
+          height = 1024,
+          guidance_scale = 5.0,
           num_inference_steps = 25
           ):
     if randomize_seed:
@@ -86,15 +97,15 @@ def infer(prompt,
     if ip_adapter_image is None:
         pipe_t2i.to(device)
         image = pipe_t2i(
-            prompt = prompt, 
+            prompt = prompt,
             negative_prompt = negative_prompt,
-            guidance_scale = guidance_scale, 
-            num_inference_steps = num_inference_steps, 
-            width = width, 
+            guidance_scale = guidance_scale,
+            num_inference_steps = num_inference_steps,
+            width = width,
             height = height,
             generator = generator
-        ).images[0] 
-        image.save("generated_image.jpg")
+        ).images[0]
+        image.save("generated_image.jpg")  # 파일 확장자를 .jpg로 변경
         return image, "generated_image.jpg"
     else:
         pipe_i2i.to(device)
@@ -104,24 +115,54 @@ def infer(prompt,
         image = pipe_i2i(
                 prompt=prompt,
                 ip_adapter_image=[ip_adapter_image],
-                negative_prompt=negative_prompt, 
+                negative_prompt=negative_prompt,
                 height=height,
                 width=width,
-                num_inference_steps=num_inference_steps, 
+                num_inference_steps=num_inference_steps,
                 guidance_scale=guidance_scale,
                 num_images_per_prompt=1,
                 generator=generator
             ).images[0]
-        image.save("generated_image.jpg")
+        image.save("generated_image.jpg")  # 파일 확장자를 .jpg로 변경
         return image, "generated_image.jpg"
 
+# 사진 설명 기능 추가를 위한 참조 코드 통합
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to("cuda").eval()
+processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+
+def modify_caption(caption: str) -> str:
+    prefix_substrings = [
+        ('captured from ', ''),
+        ('captured at ', '')
+    ]
+    pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
+    replacers = {opening.lower(): replacer for opening, replacer in prefix_substrings}
+    def replace_fn(match):
+        return replacers[match.group(0).lower()]
+    modified_caption = re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
+    return modified_caption if modified_caption != caption else caption
+    
+@spaces.GPU
 def describe_image(image):
-    # CPU에서 전처리
-    inputs = clip_processor(images=image, return_tensors="pt").to("cpu")
-    with torch.no_grad():
-        outputs = clip_model.get_image_features(**inputs).to("cpu")
-    # 여기에 적절한 텍스트 설명 생성을 위한 추가 로직이 필요할 수 있음
-    return "이미지에 대한 설명"
+    image = Image.fromarray(image)
+    task_prompt = "<DESCRIPTION>"
+    prompt = task_prompt + "Describe this image in great detail."
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+    return modify_caption(parsed_answer["<DESCRIPTION>"])
 
 css="""
 #col-left {
@@ -137,6 +178,21 @@ css="""
 with gr.Blocks(css=css) as Kolors:
     with gr.Row():
         with gr.Column(elem_id="col-left"):
+            with gr.Row():
+                korean_prompt = gr.Textbox(
+                    label="한국어 프롬프트 입력",
+                    placeholder="한국어로 원하는 프롬프트를 입력하세요",
+                    lines=2
+                )
+            with gr.Row():
+                generate_prompt_button = gr.Button("Generate Prompt")
+            with gr.Row():
+                optimized_prompt = gr.Textbox(
+                    label="최적화된 프롬프트 생성",
+                    placeholder=" ",
+                    lines=2,
+                    interactive=False
+                )
             with gr.Row():
                 generated_prompt = gr.Textbox(
                     label="프롬프트 입력",
@@ -205,16 +261,24 @@ with gr.Blocks(css=css) as Kolors:
             result = gr.Image(label="Result", show_label=False)
             download_button = gr.File(label="Download Image")
             image_description = gr.Textbox(label="Image Description", placeholder="이미지 분석 결과가 여기에 표시됩니다.", interactive=False)
+            analyze_button = gr.Button("Analyze Image")
 
+    # 최적화된 프롬프트 생성 및 결과 표시
+    generate_prompt_button.click(
+        fn=generate_prompt,
+        inputs=[korean_prompt],
+        outputs=[optimized_prompt]
+    )
+    
     # 이미지 생성 및 다운로드 파일 경로 설정
     run_button.click(
         fn=infer,
         inputs=[generated_prompt, ip_adapter_image, ip_adapter_scale, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs=[result, download_button]
     )
-
+    
     # 이미지 설명 생성
-    ip_adapter_image.change(
+    analyze_button.click(
         fn=describe_image,
         inputs=[ip_adapter_image],
         outputs=[image_description]