Update app.py

app.py

@@ -34,13 +34,14 @@ prompt_values = df.values.flatten()
 with open('loras.json', 'r') as f:
     loras = json.load(f)
 
-# Initialize the base model
-dtype = torch.bfloat16
-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="auto")
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, device_map="balanced")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, device_map="balanced").to(device)
 
 # LoRA를 위한 설정
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
@@ -55,12 +56,8 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     tokenizer=pipe.tokenizer,
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype
-)
-
-# Upscale을 위한 ControlNet 설정
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+    torch_dtype=dtype,
+    device_map="balanced"
 ).to(device)
 
 # Upscale 파이프라인 설정 (기존 pipe 재사용)
@@ -72,9 +69,9 @@ pipe_upscale = FluxControlNetPipeline(
     scheduler=pipe.scheduler,
     safety_checker=pipe.safety_checker,
     feature_extractor=pipe.feature_extractor,
-    controlnet=controlnet
-)
-
+    controlnet=controlnet,
+    device_map="balanced"
+).to(device)