Update app.py

app.py

@@ -18,47 +18,32 @@ from refiners.fluxion.utils import no_grad
 from refiners.solutions import BoxSegmenter
 from transformers import GroundingDinoForObjectDetection, GroundingDinoProcessor
 from diffusers import FluxPipeline
-# 상단에 import 추가
-# 상단에 import 추가
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-
 import gc
-import torch.cuda.amp as amp
 
 def clear_memory():
     """메모리 정리 함수"""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
     gc.collect()
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.empty_cache()
+        except:
+            pass
 
-# GPU 메모리 관리 설정
-if torch.cuda.is_available():
-    torch.cuda.empty_cache()
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.matmul.allow_tf32 = True
-    
-    # 메모리 분할 설정
-    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = (
-        "max_split_size_mb:128,"
-        "garbage_collection_threshold:0.8,"
-        "memory_fraction:0.9"
-    )
-
-    
-# 자동 혼합 정밀도(Automatic Mixed Precision) 설정
+# GPU 설정을 try-except로 감싸기
 if torch.cuda.is_available():
-    scaler = torch.amp.GradScaler('cuda')
-else:
-    scaler = None
-
-
+    try:
+        torch.cuda.empty_cache()
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+    except:
+        print("Warning: Could not configure CUDA settings")
 
+# 번역 모델 초기화
 model_name = "Helsinki-NLP/opus-mt-ko-en"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cpu')  # CPU로 강제 지정
-translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)  # 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:
     """한글 텍스트를 영어로 번역"""
@@ -72,8 +57,6 @@ def translate_to_english(text: str) -> str:
         print(f"Translation error: {str(e)}")
         return text
 
-
-
 BoundingBox = tuple[int, int, int, int]
 
 pillow_heif.register_heif_opener()
@@ -102,15 +85,13 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
-
-# FLUX 파이프라인 초기화
 # FLUX 파이프라인 초기화
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.float16,  # A100에 최적화된 float16 사용
+    torch_dtype=torch.float16,
     use_auth_token=HF_TOKEN
 )
-pipe.enable_attention_slicing(slice_size="auto")  # 메모리 사용량 최적화
+pipe.enable_attention_slicing(slice_size="auto")
 
 # LoRA 가중치 로드
 pipe.load_lora_weights(
@@ -122,16 +103,12 @@ pipe.load_lora_weights(
 )
 pipe.fuse_lora(lora_scale=0.125)
 
-# GPU 메모리 최적화
-if torch.cuda.is_available():
-    pipe.to("cuda")
-    # 메모리 최적화를 위한 추가 설정
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.matmul.allow_tf32 = True
-
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 단일 GPU 사용
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"  # CUDA 메모리 할당 설정
-
+# GPU 설정을 try-except로 감싸기
+try:
+    if torch.cuda.is_available():
+        pipe.to("cuda")
+except:
+    print("Warning: Could not move pipeline to CUDA")
 
 class timer:
     def __init__(self, method_name="timed process"):
@@ -210,7 +187,6 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
         width, height = calculate_dimensions(aspect_ratio)
         width, height = adjust_size_to_multiple_of_8(width, height)
         
-        # A100 메모리 제한을 고려한 최대 크기 설정
         max_size = 768
         if width > max_size or height > max_size:
             ratio = max_size / max(width, height)
@@ -218,24 +194,24 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             height = int(height * ratio)
             width, height = adjust_size_to_multiple_of_8(width, height)
 
-        clear_memory()  # 생성 전 메모리 정리
-        
         with timer("Background generation"):
-            with torch.inference_mode(), torch.cuda.amp.autocast():
-                image = pipe(
-                    prompt=prompt,
-                    width=width,
-                    height=height,
-                    num_inference_steps=8,
-                    guidance_scale=4.0,
-                    max_length=77,
-                ).images[0]
-        
-        clear_memory()  # 생성 후 메모리 정리
+            try:
+                with torch.inference_mode():
+                    image = pipe(
+                        prompt=prompt,
+                        width=width,
+                        height=height,
+                        num_inference_steps=8,
+                        guidance_scale=4.0,
+                        max_length=77,
+                    ).images[0]
+            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)}")
-        clear_memory()  # 오류 발생 시에도 메모리 정리
         return Image.new('RGB', (512, 512), 'white')
 
 
@@ -355,21 +331,18 @@ 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]:
     try:
-        clear_memory()  # 처리 전 메모리 정리        
         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}")
         
         try:
-            # 프롬프트 번역 시도
             prompt = translate_to_english(prompt)
             if bg_prompt:
                 bg_prompt = translate_to_english(bg_prompt)
         except Exception as e:
             print(f"Translation error (continuing with original text): {str(e)}")
         
-        # Process the image
         results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
         
         if bg_prompt:
@@ -390,9 +363,8 @@ def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None,
     except Exception as e:
         print(f"Error in process_prompt: {str(e)}")
         raise gr.Error(str(e))
-
     finally:
-        clear_memory()  # 처리 후 메모리 정리
+        clear_memory()
         
 def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.Image]:
     try: