Update app.py

app.py

@@ -2,7 +2,7 @@ import spaces
 import random
 import torch
 from huggingface_hub import snapshot_download
-from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, CLIPModel, CLIPTokenizer
 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,7 +10,6 @@ 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
 
 device = "cuda"
@@ -27,6 +26,11 @@ image_encoder = CLIPVisionModelWithProjection.from_pretrained(f'{ckpt_IPA_dir}/i
 ip_img_size = 336
 clip_image_processor = CLIPImageProcessor(size=ip_img_size, crop_size=ip_img_size)
 
+# CLIP 모델 및 토크나이저 로드
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+clip_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
 pipe_t2i = pipeline_stable_diffusion_xl_chatglm_256.StableDiffusionXLPipeline(
     vae=vae,
     text_encoder=text_encoder, 
@@ -103,6 +107,13 @@ def infer(prompt,
         image.save("generated_image.jpg")  # 파일 확장자를 .jpg로 변경
         return image, "generated_image.jpg"
 
+def describe_image(image):
+    image = clip_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        text_features = clip_model.get_image_features(**image)
+    text = clip_tokenizer.decode(torch.argmax(text_features, dim=-1))
+    return text
+
 css="""
 #col-left {
     margin: 0 auto;
@@ -184,6 +195,7 @@ with gr.Blocks(css=css) as Kolors:
         with gr.Column(elem_id="col-right"):
             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)
 
     # 이미지 생성 및 다운로드 파일 경로 설정
     run_button.click(
@@ -192,4 +204,11 @@ with gr.Blocks(css=css) as Kolors:
         outputs=[result, download_button]
     )
 
+    # 이미지 설명 생성
+    ip_adapter_image.change(
+        fn=describe_image,
+        inputs=[ip_adapter_image],
+        outputs=[image_description]
+    )
+
 Kolors.queue().launch(debug=True)