Update app.py

app.py

@@ -3,6 +3,7 @@ import time
 from collections.abc import Sequence
 from typing import Any, cast
 import os
+import gc
 from huggingface_hub import login, hf_hub_download
 
 import gradio as gr
@@ -19,35 +20,36 @@ from refiners.solutions import BoxSegmenter
 from transformers import GroundingDinoForObjectDetection, GroundingDinoProcessor
 from diffusers import FluxPipeline
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-import gc
 
+#############################################################
+# 메모리 정리 함수
 def clear_memory():
-    """메모리 정리 함수"""
     gc.collect()
     try:
         if torch.cuda.is_available():
             with torch.cuda.device(0):  # 명시적으로 device 0 사용
                 torch.cuda.empty_cache()
-    except:
+    except Exception as e:
         pass
 
-# GPU 설정
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명시적으로 cuda:0 지정
-
-# GPU 설정을 try-except로 감싸기
+#############################################################
+# GPU 설정 (Zero GPU 환경)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
     try:
         with torch.cuda.device(0):
             torch.cuda.empty_cache()
             torch.backends.cudnn.benchmark = True
             torch.backends.cuda.matmul.allow_tf32 = True
-    except:
+    except Exception as e:
         print("Warning: Could not configure CUDA settings")
 
-# 번역 모델 초기화
+#############################################################
+# 번역 모델 초기화 (CPU에서 동작)
 model_name = "Helsinki-NLP/opus-mt-ko-en"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cpu')
+# 번역 모델은 CPU에 올림
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to("cpu")
 translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)
 
 def translate_to_english(text: str) -> str:
@@ -67,6 +69,7 @@ BoundingBox = tuple[int, int, int, int]
 pillow_heif.register_heif_opener()
 pillow_heif.register_avif_opener()
 
+#############################################################
 # HF 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
@@ -77,7 +80,8 @@ try:
 except Exception as e:
     raise ValueError(f"Failed to login to Hugging Face: {str(e)}")
 
-# 모델 초기화
+#############################################################
+# 객체 분할 모델 초기화
 segmenter = BoxSegmenter(device="cpu")
 segmenter.device = device
 segmenter.model = segmenter.model.to(device=segmenter.device)
@@ -88,15 +92,14 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
-# FLUX 파이프라인 초기화
+#############################################################
+# FLUX 파이프라인 초기화 (Zero GPU용)
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     torch_dtype=torch.float16,
     use_auth_token=HF_TOKEN
 )
 pipe.enable_attention_slicing(slice_size="auto")
-
-# LoRA 가중치 로드
 pipe.load_lora_weights(
     hf_hub_download(
         "ByteDance/Hyper-SD",
@@ -105,14 +108,14 @@ pipe.load_lora_weights(
     )
 )
 pipe.fuse_lora(lora_scale=0.125)
-
-# GPU 설정을 try-except로 감싸기
 try:
     if torch.cuda.is_available():
-        pipe = pipe.to("cuda:0")  # 명시적으로 cuda:0 지정
+        pipe = pipe.to("cuda:0")  # 명시적으로 cuda:0로 이동
 except Exception as e:
     print(f"Warning: Could not move pipeline to CUDA: {str(e)}")
 
+#############################################################
+# 타이머 클래스
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name
@@ -123,6 +126,8 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
+#############################################################
+# 유틸리티 함수들
 def bbox_union(bboxes: Sequence[list[int]]) -> BoundingBox | None:
     if not bboxes:
         return None
@@ -166,15 +171,12 @@ def apply_mask(img: Image.Image, mask_img: Image.Image, defringe: bool = True) -
     result.paste(img, (0, 0), mask_img)
     return result
 
-
 def adjust_size_to_multiple_of_8(width: int, height: int) -> tuple[int, int]:
-    """이미지 크기를 8의 배수로 조정하는 함수"""
     new_width = ((width + 7) // 8) * 8
     new_height = ((height + 7) // 8) * 8
     return new_width, new_height
 
 def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int, int]:
-    """선택된 비율에 따라 이미지 크기 계산"""
     if aspect_ratio == "1:1":
         return base_size, base_size
     elif aspect_ratio == "16:9":
@@ -185,7 +187,9 @@ def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int,
         return base_size * 4 // 3, base_size
     return base_size, base_size
 
-@spaces.GPU(duration=20)  # 40초에서 20초로 감소
+#############################################################
+# 배경 생성 함수 (Zero GPU에 맞게 수정)
+@spaces.GPU(duration=20)
 def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
     try:
         width, height = calculate_dimensions(aspect_ratio)
@@ -197,7 +201,7 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             width = int(width * ratio)
             height = int(height * ratio)
             width, height = adjust_size_to_multiple_of_8(width, height)
-
+        
         with timer("Background generation"):
             try:
                 with torch.inference_mode():
@@ -211,7 +215,6 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             except Exception as e:
                 print(f"Pipeline error: {str(e)}")
                 return Image.new('RGB', (width, height), 'white')
-                
         return image
     except Exception as e:
         print(f"Background generation error: {str(e)}")
@@ -233,7 +236,6 @@ def create_position_grid():
     """
 
 def calculate_object_position(position: str, bg_size: tuple[int, int], obj_size: tuple[int, int]) -> tuple[int, int]:
-    """오브젝트의 위치 계산"""
     bg_width, bg_height = bg_size
     obj_width, obj_height = obj_size
     
@@ -252,28 +254,21 @@ def calculate_object_position(position: str, bg_size: tuple[int, int], obj_size:
     return positions.get(position, positions["bottom-center"])
 
 def resize_object(image: Image.Image, scale_percent: float) -> Image.Image:
-    """오브젝트 크기 조정"""
     width = int(image.width * scale_percent / 100)
     height = int(image.height * scale_percent / 100)
     return image.resize((width, height), Image.Resampling.LANCZOS)
 
 def combine_with_background(foreground: Image.Image, background: Image.Image, 
-                          position: str = "bottom-center", scale_percent: float = 100) -> Image.Image:
-    """전경과 배경 합성 함수"""
-    # 배경 이미지 준비
+                            position: str = "bottom-center", scale_percent: float = 100) -> Image.Image:
     result = background.convert('RGBA')
-    
-    # 오브젝트 크기 조정
     scaled_foreground = resize_object(foreground, scale_percent)
-    
-    # 오브젝트 위치 계산
     x, y = calculate_object_position(position, result.size, scaled_foreground.size)
-    
-    # 합성
     result.paste(scaled_foreground, (x, y), scaled_foreground)
     return result
 
-@spaces.GPU(duration=30)  # 120초에서 30초로 감소
+#############################################################
+# GPU 처리 함수 (Zero GPU에 맞게 수정)
+@spaces.GPU(duration=30)
 def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Image.Image, BoundingBox | None, list[str]]:
     time_log: list[str] = []
     try:
@@ -294,6 +289,8 @@ def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Im
         print(f"GPU process error: {str(e)}")
         raise
 
+#############################################################
+# 전체 처리 함수
 def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str | None = None, aspect_ratio: str = "1:1") -> tuple[tuple[Image.Image, Image.Image, Image.Image], gr.DownloadButton]:
     try:
         # 입력 이미지 크기 제한
@@ -302,8 +299,7 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
             ratio = max_size / max(img.width, img.height)
             new_size = (int(img.width * ratio), int(img.height * ratio))
             img = img.resize(new_size, Image.LANCZOS)
-
-        # CUDA 메모리 관리 수정
+        
         try:
             if torch.cuda.is_available():
                 current_device = torch.cuda.current_device()
@@ -311,19 +307,19 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
                     torch.cuda.empty_cache()
         except Exception as e:
             print(f"CUDA memory management failed: {e}")
-
+        
         with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
             mask, bbox, time_log = _gpu_process(img, prompt)
             masked_alpha = apply_mask(img, mask, defringe=True)
-
+        
         if bg_prompt:
             background = generate_background(bg_prompt, aspect_ratio)
             combined = background
         else:
             combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
-
+        
         clear_memory()
-
+        
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp:
             combined.save(temp.name)
             return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)
@@ -335,15 +331,12 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
 def on_change_bbox(prompts: dict[str, Any] | None):
     return gr.update(interactive=prompts is not None)
 
-
 def on_change_prompt(img: Image.Image | None, prompt: str | None, bg_prompt: str | None = None):
     return gr.update(interactive=bool(img and prompt))
 
-
-        
 def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None, 
-                  aspect_ratio: str = "1:1", position: str = "bottom-center", 
-                  scale_percent: float = 100) -> tuple[Image.Image, Image.Image]:
+                   aspect_ratio: str = "1:1", position: str = "bottom-center", 
+                   scale_percent: float = 100) -> tuple[Image.Image, Image.Image]:
     try:
         if img is None or prompt.strip() == "":
             raise gr.Error("Please provide both image and prompt")
@@ -379,7 +372,7 @@ def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None,
         raise gr.Error(str(e))
     finally:
         clear_memory()
-        
+
 def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.Image]:
     try:
         if img is None or box_input.strip() == "":
@@ -393,15 +386,11 @@ def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.I
         except:
             raise gr.Error("Invalid box format. Please provide [xmin, ymin, xmax, ymax]")
         
-        # Process the image
         results, _ = _process(img, bbox)
-        
-        # 합성된 이미지와 추출된 이미지만 반환
         return results[1], results[2]
     except Exception as e:
         raise gr.Error(str(e))
 
-# Event handler functions 수정
 def update_process_button(img, prompt):
     return gr.update(
         interactive=bool(img and prompt),
@@ -418,7 +407,7 @@ def update_box_button(img, box_input):
     except:
         return gr.update(interactive=False, variant="secondary")
 
-
+#############################################################
 # CSS 정의
 css = """
 footer {display: none}
@@ -482,9 +471,7 @@ button.primary:hover {
 }
 """
 
-# UI 구성
-# UI 구성 부분에서 process_btn을 위로 이동하고 position_grid.click 부분 제거
-
+#############################################################
 # UI 구성
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.HTML("""
@@ -493,7 +480,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             <p>AI Integrated Image Creator: Extract objects, generate backgrounds, and adjust ratios and positions to create complete images with AI.</p>
         </div>
     """)
-
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(
@@ -521,7 +507,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     visible=True,
                     scale=1
                 )
-
             with gr.Row(visible=False) as object_controls:
                 with gr.Column(scale=1):
                     with gr.Row():
@@ -545,17 +530,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                         step=5,
                         label="Object Size (%)"
                     )
-
             process_btn = gr.Button(
                 "Process",
                 variant="primary",
                 interactive=False
             )
-
             # 각 버튼에 대한 클릭 이벤트 처리
             def update_position(new_position):
                 return new_position
-
             btn_top_left.click(fn=lambda: update_position("top-left"), outputs=position)
             btn_top_center.click(fn=lambda: update_position("top-center"), outputs=position)
             btn_top_right.click(fn=lambda: update_position("top-right"), outputs=position)
@@ -565,7 +547,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             btn_bottom_left.click(fn=lambda: update_position("bottom-left"), outputs=position)
             btn_bottom_center.click(fn=lambda: update_position("bottom-center"), outputs=position)
             btn_bottom_right.click(fn=lambda: update_position("bottom-right"), outputs=position)
-
         with gr.Column(scale=1):
             with gr.Row():
                 combined_image = gr.Image(
@@ -581,7 +562,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     type="pil",
                     height=256
                 )
-
     # Event bindings
     input_image.change(
         fn=update_process_button,
@@ -589,29 +569,24 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=process_btn,
         queue=False
     )
-    
     text_prompt.change(
         fn=update_process_button,
         inputs=[input_image, text_prompt],
         outputs=process_btn,
         queue=False
     )
-
     def update_controls(bg_prompt):
-        """배경 프롬프트 입력 여부에 따라 컨트롤 표시 업데이트"""
         is_visible = bool(bg_prompt)
         return [
-            gr.update(visible=is_visible),  # aspect_ratio
-            gr.update(visible=is_visible),  # object_controls
+            gr.update(visible=is_visible),
+            gr.update(visible=is_visible),
         ]
-
     bg_prompt.change(
         fn=update_controls,
         inputs=bg_prompt,
         outputs=[aspect_ratio, object_controls],
         queue=False
     )
-
     process_btn.click(
         fn=process_prompt,
         inputs=[
@@ -625,12 +600,23 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=[combined_image, extracted_image],
         queue=True
     )
-
-
-demo.queue(max_size=5)  # 큐 크기 제한
+    # 예제 섹션 추가
+    with gr.Accordion("Show Example", open=True):
+        gr.Markdown("### Example")
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("**Upload Image(aa1.png)**")
+                gr.Image(value="aa1.png", label="Upload")
+            with gr.Column():
+                gr.Markdown("**Cut Object (aa2.png)**<br>(Prompt: 'text')", elem_classes="center")
+                gr.Image(value="aa2.png", label="Object")
+            with gr.Column():
+                gr.Markdown("**Generated Image (aa3.png)**<br>(Background Prompt: 'alps mountain')", elem_classes="center")
+                gr.Image(value="aa3.png", label="Output")
+demo.queue(max_size=5)
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
     share=False,
-    max_threads=2  # 스레드 수 제한
+    max_threads=2
 )