Update app.py

app.py

@@ -71,19 +71,23 @@ def translate_to_english(text: str) -> str:
         print(f"Translation error: {str(e)}")
         return text
         
+# FLUX 파이프라인 초기화 부분 수정
 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  # 안전성 검사 비활성화
+        safety_checker=None,
+        variant="fp16",  # fp16 변형 사용
+        device_map="auto"  # 자동 장치 매핑
     )
     print("FLUX pipeline initialized successfully")
     
-    # 메모리 최적화 설정
-    pipe.enable_attention_slicing(slice_size="auto")
+    # 메모리 최적화 설정 강화
+    pipe.enable_attention_slicing(slice_size=1)  # 더 작은 슬라이스 크기
     pipe.enable_model_cpu_offload()  # CPU 오프로딩 활성화
+    pipe.enable_sequential_cpu_offload()  # 순차적 CPU 오프로딩
     print("Pipeline optimization settings applied")
 
 except Exception as e:
@@ -93,34 +97,28 @@ except Exception as e:
 # LoRA 가중치 로드 부분 수정
 print("Loading LoRA weights...")
 try:
-    # LoRA 파일 경로 확인
     lora_path = hf_hub_download(
         repo_id="openfree/myt-flux-fantasy",
-        filename="myt-flux-fantasy.safetensors",  # 정확한 파일명 확인 필요
+        filename="myt-flux-fantasy.safetensors",
         use_auth_token=HF_TOKEN
     )
     print(f"LoRA weights downloaded to: {lora_path}")
     
-    # LoRA 가중치 로드
-    pipe.load_lora_weights(lora_path)
+    # LoRA 가중치 로드 (메모리 효율적 방식)
+    pipe.load_lora_weights(lora_path, adapter_name="fantasy")
     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)}")
-    print(f"Full error details: {repr(e)}")
-    raise ValueError("Failed to load LoRA weights. Please check your HF_TOKEN and model access.")
+    raise ValueError("Failed to load LoRA weights")
+
 
-# GPU 이동 부분 수정
-if torch.cuda.is_available():
-    try:
-        print("Moving pipeline to GPU...")
-        pipe = pipe.to("cuda:0")
-        print("Pipeline successfully moved to GPU")
-        print(f"Current device: {pipe.device}")
-    except Exception as e:
-        print(f"Warning: Could not move pipeline to CUDA: {str(e)}")
-        print("Falling back to CPU")
 
 
 # 저장 디렉토리 설정
@@ -139,6 +137,7 @@ def save_generated_image(image, prompt):
     image.save(filepath)
     return filepath
 
+# generate_image 함수 수정
 @spaces.GPU(duration=60)
 def generate_image(
     prompt: str,
@@ -153,18 +152,19 @@ def generate_image(
     try:
         print(f"\nStarting image generation with prompt: {prompt}")
         
-        # 프롬프트 번역
+        # 메모리 정리
+        clear_memory()
+        
         translated_prompt = translate_to_english(prompt)
         print(f"Translated prompt: {translated_prompt}")
         
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-        print(f"Using seed: {seed}")
         
         generator = torch.Generator(device=device).manual_seed(seed)
         
-        print("Starting inference...")
-        with torch.inference_mode():
+        # 배치 크기 1로 고정하여 메모리 사용 최소화
+        with torch.inference_mode(), torch.cuda.amp.autocast():
             image = pipe(
                 prompt=translated_prompt,
                 width=width,
@@ -172,20 +172,20 @@ def generate_image(
                 num_inference_steps=num_inference_steps,
                 guidance_scale=guidance_scale,
                 generator=generator,
+                num_images_per_prompt=1,
             ).images[0]
-        print("Image generation completed successfully")
         
         filepath = save_generated_image(image, translated_prompt)
-        print(f"Image saved to: {filepath}")
+        
+        # 메모리 정리
+        clear_memory()
         
         return image, seed
         
     except Exception as e:
         print(f"Error in generate_image: {str(e)}")
-        print(f"Full error details: {repr(e)}")
-        raise gr.Error(f"Image generation failed: {str(e)}")
-    finally:
         clear_memory()
+        raise gr.Error(f"Image generation failed: {str(e)}")
 
 def add_text_with_stroke(draw, text, x, y, font, text_color, stroke_width):
     """텍스트에 외곽선을 추가하는 함수"""