Update app.py

app.py

@@ -76,25 +76,29 @@ print("Initializing FLUX pipeline...")
 try:
     pipe = FluxPipeline.from_pretrained(
         "black-forest-labs/FLUX.1-dev",
-        torch_dtype=torch.float16,  # 반정밀도 사용
+        torch_dtype=torch.float16,
         use_auth_token=HF_TOKEN,
         safety_checker=None,
-        variant="fp16",  # fp16 변형 사용
-        device_map="auto"  # 자동 장치 매핑
+        device_map="balanced"  # 'auto' 대신 'balanced' 사용
     )
     print("FLUX pipeline initialized successfully")
     
-    # 메모리 최적화 설정 강화
-    pipe.enable_attention_slicing(slice_size=1)  # 더 작은 슬라이스 크기
-    pipe.enable_model_cpu_offload()  # CPU 오프로딩 활성화
-    pipe.enable_sequential_cpu_offload()  # 순차적 CPU 오프로딩
+    # 메모리 최적화 설정
+    pipe.enable_attention_slicing(slice_size=1)
+    pipe.enable_model_cpu_offload()
     print("Pipeline optimization settings applied")
 
+    # 추가 메모리 최적화
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.matmul.allow_tf32 = True
+
 except Exception as e:
     print(f"Error initializing FLUX pipeline: {str(e)}")
     raise
 
-# LoRA 가중치 로드 부분 수정
+# LoRA 가중치 로드 부분
 print("Loading LoRA weights...")
 try:
     lora_path = hf_hub_download(
@@ -104,39 +108,14 @@ try:
     )
     print(f"LoRA weights downloaded to: {lora_path}")
     
-    # LoRA 가중치 로드 (메모리 효율적 방식)
-    pipe.load_lora_weights(lora_path, adapter_name="fantasy")
+    pipe.load_lora_weights(lora_path)
     pipe.fuse_lora(lora_scale=0.125)
-    
-    # 불필요한 메모리 정리
-    torch.cuda.empty_cache()
-    gc.collect()
-    
     print("LoRA weights loaded and fused successfully")
 
 except Exception as e:
     print(f"Error loading LoRA weights: {str(e)}")
     raise ValueError("Failed to load LoRA weights")
 
-
-
-
-# 저장 디렉토리 설정
-SAVE_DIR = "saved_images"
-if not os.path.exists(SAVE_DIR):
-    os.makedirs(SAVE_DIR, exist_ok=True)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def save_generated_image(image, prompt):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    unique_id = str(uuid.uuid4())[:8]
-    filename = f"{timestamp}_{unique_id}.png"
-    filepath = os.path.join(SAVE_DIR, filename)
-    image.save(filepath)
-    return filepath
-
 # generate_image 함수 수정
 @spaces.GPU(duration=60)
 def generate_image(
@@ -150,21 +129,17 @@ def generate_image(
     progress: gr.Progress = gr.Progress()
 ):
     try:
-        print(f"\nStarting image generation with prompt: {prompt}")
-        
-        # 메모리 정리
         clear_memory()
         
         translated_prompt = translate_to_english(prompt)
-        print(f"Translated prompt: {translated_prompt}")
+        print(f"Processing prompt: {translated_prompt}")
         
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
         generator = torch.Generator(device=device).manual_seed(seed)
         
-        # 배치 크기 1로 고정하여 메모리 사용 최소화
-        with torch.inference_mode(), torch.cuda.amp.autocast():
+        with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True):
             image = pipe(
                 prompt=translated_prompt,
                 width=width,
@@ -176,16 +151,34 @@ def generate_image(
             ).images[0]
         
         filepath = save_generated_image(image, translated_prompt)
-        
-        # 메모리 정리
-        clear_memory()
-        
         return image, seed
         
     except Exception as e:
-        print(f"Error in generate_image: {str(e)}")
-        clear_memory()
+        print(f"Generation error: {str(e)}")
         raise gr.Error(f"Image generation failed: {str(e)}")
+    finally:
+        clear_memory()
+
+
+
+
+# 저장 디렉토리 설정
+SAVE_DIR = "saved_images"
+if not os.path.exists(SAVE_DIR):
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+def save_generated_image(image, prompt):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"{timestamp}_{unique_id}.png"
+    filepath = os.path.join(SAVE_DIR, filename)
+    image.save(filepath)
+    return filepath
+
+
 
 def add_text_with_stroke(draw, text, x, y, font, text_color, stroke_width):
     """텍스트에 외곽선을 추가하는 함수"""