Update app.py

app.py

@@ -54,22 +54,26 @@ model = model.cuda()
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
 # FLUX 모델 설정
-dtype = torch.float16  # bfloat16 대신 float16 사용
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"  # 초기에 CPU에 로드
+dtype = torch.float32  # float32 사용
 
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype)
-pipe = pipe.to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, use_fast=False)
 
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 2048
+# 메모리 사용량 최적화
+if torch.cuda.is_available():
+    pipe.enable_sequential_cpu_offload()
+    pipe.enable_attention_slicing(1)
 
 @spaces.GPU(duration=300)
 def infer_t2i(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    with torch.cuda.amp.autocast():
+    generator = torch.Generator().manual_seed(seed)
+    
+    # 추론 시 GPU 사용 (가능한 경우)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    with torch.no_grad():
         image = pipe(
             prompt=prompt, 
             width=width,
@@ -78,11 +82,10 @@ def infer_t2i(prompt, seed=42, randomize_seed=False, width=1024, height=1024, gu
             generator=generator,
             guidance_scale=guidance_scale
         ).images[0]
+    
     torch.cuda.empty_cache()
     return image, seed
-
-
-
+    
 @spaces.GPU(duration=300)
 def infer(image, prompt, steps=50, cfg_scale=7.5, eta=1.0, fs=3, seed=123, video_length=2):
     # 한글 입력 감지 및 번역