update minor

.gitignore

@@ -0,0 +1,4 @@
+models
+*.log
+*.nohup
+*.db

README.md

@@ -8,6 +8,13 @@ app_file: gradio_imagecompletion.py
 pinned: false
 ---
 
+# Note
+
+- 주어진 이미지의 하단부 절반을 지우고 새로 채워서 그려주는 AI 데모입니다
+- ImageGPT 활용
+    - Paper: https://arxiv.org/abs/2109.10282
+    - Code: https://huggingface.co/spaces/nielsr/imagegpt-completion
+
 # Configuration
 
 `title`: _string_  

gradio_imagecompletion.py

@@ -13,7 +13,6 @@ import matplotlib.pyplot as plt
 
 import os
 import numpy as np
-import requests
 from glob import glob
 import gradio as gr
 from loguru import logger
@@ -30,13 +29,20 @@ logger.add('app.log', mode='a')
 logger.info('===== APP RESTARTED =====')
 
 # ========== Models ==========
+# MODEL_DIR = 'models'
+# os.environ['TORCH_HOME'] = MODEL_DIR
+# os.environ['TF_HOME'] = MODEL_DIR
 feature_extractor = ImageGPTFeatureExtractor.from_pretrained(
-    "openai/imagegpt-medium")
+    "openai/imagegpt-medium", 
+    # cache_dir=MODEL_DIR
+    )
 model = ImageGPTForCausalImageModeling.from_pretrained(
-    "openai/imagegpt-medium")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
+    "openai/imagegpt-medium", 
+    # cache_dir=MODEL_DIR
+    )
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(DEVICE)
+logger.info(f'model loaded (DEVICE:{DEVICE})')
 
 def process_image(image):
     logger.info('--- image file received')
@@ -56,7 +62,7 @@ def process_image(image):
     
     # generate (no beam search)
     context = np.concatenate((np.full((batch_size, 1), model.config.vocab_size - 1), primers), axis=1)
-    context = torch.tensor(context).to(device)
+    context = torch.tensor(context).to(DEVICE)
     output = model.generate(input_ids=context, max_length=n_px*n_px + 1, temperature=1.0, do_sample=True, top_k=40)
     # decode back to images (convert color cluster tokens back to pixels)
     samples = output[:,1:].cpu().detach().numpy()
@@ -71,13 +77,14 @@ def process_image(image):
     
     # return as PIL Image
     completion = Image.fromarray(result)
+    logger.info('--- image generated')
     return completion
 
 
 iface = gr.Interface(
     process_image,
     title="이미지의 절반을 지우고 절반을 채워 넣어주는 Image Completion 데모입니다 (ImageGPT)",
-    description='주어진 이미지의 절반 아래를 AI가 채워 넣어줍니다',
+    description='주어진 이미지의 절반 아래를 AI가 채워 넣어줍니다 (CPU로 약 100초 정도 소요됩니다)',
     inputs=gr.inputs.Image(type="pil", label='인풋 이미지'),
     outputs=gr.outputs.Image(type="pil", label='AI가 그린 결과'),
     examples=examples,