Update app.py

app.py

@@ -24,7 +24,7 @@ import warnings
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
 # 번역 모델 로드
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device=0 if torch.cuda.is_available() else -1)
 
 #Load prompts for randomization
 df = pd.read_csv('prompts.csv', header=None)
@@ -40,12 +40,12 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 공통 FLUX 모델 로드
 base_model = "black-forest-labs/FLUX.1-dev"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, device_map="balanced").to(device)
-
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype)
+pipe.to(device)  # 여기서 한 번만 device로 이동
 
 # LoRA를 위한 설정
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype)
 
 # Image-to-Image 파이프라인 설정
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
@@ -56,14 +56,12 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     tokenizer=pipe.tokenizer,
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype,
-    device_map="balanced"
-).to(device)
+    torch_dtype=dtype
+)
 
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-).to(device)
-
+)
 
 # Upscale 파이프라인 설정 (기존 pipe 재사용)
 pipe_upscale = FluxControlNetPipeline(
@@ -74,9 +72,13 @@ pipe_upscale = FluxControlNetPipeline(
     scheduler=pipe.scheduler,
     safety_checker=pipe.safety_checker,
     feature_extractor=pipe.feature_extractor,
-    controlnet=controlnet,
-    device_map="balanced"
-).to(device)
+    controlnet=controlnet
+)
+
+# 모든 파이프라인을 device로 이동
+pipe_i2i.to(device)
+controlnet.to(device)
+pipe_upscale.to(device)