Update app.py

app.py

@@ -34,15 +34,15 @@ 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="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)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
@@ -60,6 +60,11 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     device_map="balanced"
 ).to(device)
 
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+).to(device)
+
+
 # Upscale 파이프라인 설정 (기존 pipe 재사용)
 pipe_upscale = FluxControlNetPipeline(
     vae=pipe.vae,