Update app.py

app.py

@@ -22,16 +22,66 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import gc
 from PIL import Image, ImageDraw, ImageFont
 
-def clear_memory():
-    """메모리 정리 함수"""
-    gc.collect()
+# 모델 초기화 부분 수정
+def initialize_models():
+    global segmenter, gd_model, gd_processor, pipe
+    
     try:
+        # GPU 메모리 정리
+        clear_memory()
+        
+        # 번역 모델은 CPU에서만 실행
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cpu')
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)
+        
+        # GroundingDINO 모델 초기화
+        gd_processor = GroundingDinoProcessor.from_pretrained(gd_model_path)
+        gd_model = GroundingDinoForObjectDetection.from_pretrained(
+            gd_model_path,
+            torch_dtype=torch.float16,  # float32 대신 float16 사용
+            device_map="auto"  # 자동 디바이스 매핑
+        )
+        
+        # Segmenter 초기화
+        segmenter = BoxSegmenter(device="cpu")
         if torch.cuda.is_available():
-            with torch.cuda.device(0):  # 명시적으로 device 0 사용
-                torch.cuda.empty_cache()
-    except:
-        pass
-
+            segmenter.to(device)
+        
+        # FLUX 파이프라인 초기화
+        pipe = FluxPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            torch_dtype=torch.float16,
+            token=HF_TOKEN
+        )
+        pipe.enable_attention_slicing(slice_size="auto")
+        
+        # LoRA 가중치 로드
+        pipe.load_lora_weights(
+            hf_hub_download(
+                "ByteDance/Hyper-SD",
+                "Hyper-FLUX.1-dev-8steps-lora.safetensors",
+                use_auth_token=HF_TOKEN
+            )
+        )
+        pipe.fuse_lora(lora_scale=0.125)
+        
+        if torch.cuda.is_available():
+            pipe = pipe.to("cuda:0")
+            
+    except Exception as e:
+        print(f"Model initialization error: {str(e)}")
+        raise
+        
+def clear_memory():
+    """메모리 정리 강화 함수"""
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    if torch.cuda.is_available():
+        with torch.cuda.device(0):
+            torch.cuda.reset_peak_memory_stats()
+            torch.cuda.empty_cache()
 # GPU 설정
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명시적으로 cuda:0 지정
 
@@ -340,51 +390,40 @@ def on_change_bbox(prompts: dict[str, Any] | 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))
 
+
+@spaces.GPU(duration=20)
 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, text_params: dict | None = None) -> tuple[Image.Image, Image.Image]:
     try:
-        if img is None or prompt.strip() == "":
-            raise gr.Error("Please provide both image and prompt")
-        
-        
-        print(f"Processing with position: {position}, scale: {scale_percent}")
+        # 입력 이미지 최적화
+        img = process_image(img)
         
-        try:
-            prompt = translate_to_english(prompt)
+        with torch.cuda.amp.autocast():  # 자동 혼합 정밀도 사용
+            results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
+            
             if bg_prompt:
-                bg_prompt = translate_to_english(bg_prompt)
-        except Exception as e:
-            print(f"Translation error (continuing with original text): {str(e)}")
-        
-        # 기존 처리 로직...
-        results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
-        
-        if bg_prompt:
-            try:
                 combined = combine_with_background(
                     foreground=results[2],
                     background=results[1],
                     position=position,
                     scale_percent=scale_percent
                 )
-            except Exception as e:
-                print(f"Combination error: {str(e)}")
-                combined = results[1]
-        else:
-            combined = results[1]
-
-        # 텍스트 추가 로직을 여기로 이동
-        if text_params and text_params.get('text'):
-            combined = add_text_to_image(combined, text_params)
-        
-        return combined, results[2]
+                
+                if text_params and text_params.get('text'):
+                    combined = add_text_to_image(combined, text_params)
+                    
+                return combined, results[2]
+            
+            return results[1], results[2]
+            
     except Exception as e:
         print(f"Error in process_prompt: {str(e)}")
         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() == "":
@@ -423,7 +462,14 @@ def update_box_button(img, box_input):
     except:
         return gr.update(interactive=False, variant="secondary")
 
-
+def process_image(img: Image.Image, max_size: int = 1024) -> Image.Image:
+    """이미지 크기 최적화"""
+    if img.width > max_size or img.height > max_size:
+        ratio = max_size / max(img.width, img.height)
+        new_size = (int(img.width * ratio), int(img.height * ratio))
+        return img.resize(new_size, Image.LANCZOS)
+    return img
+    
 # CSS 정의
 css = """
 footer {display: none}
@@ -559,9 +605,15 @@ def update_process_button(img, prompt):
         interactive=bool(img and prompt),
         variant="primary" if bool(img and prompt) else "secondary"
     )
+
+demo = gr.Blocks(
+    theme=gr.themes.Soft(),
+    css=css,
+    analytics_enabled=False,  # 분석 비활성화
+    cache_examples=False  # 예제 캐싱 비활성화
+)
     
-# UI 구성
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+
     gr.HTML("""
         <div class="main-title">
             <h1>🎨GiniGen Canvas</h1>
@@ -769,10 +821,14 @@ pipe = FluxPipeline.from_pretrained(
     token=HF_TOKEN
 )
 
-demo.queue(max_size=5)
+demo.queue(max_size=3)  # 큐 크기 감소
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
     share=False,
-    max_threads=2
+    max_threads=2,
+    enable_queue=True,
+    cache_examples=False,
+    show_error=True,
+    show_tips=False
 )