Update app.py

app.py

@@ -38,16 +38,16 @@ with open('loras.json', 'r') as 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)
-pipe.to(device)  # 여기서 한 번만 device로 이동
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
 # LoRA를 위한 설정
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype)
+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)
 
 # Image-to-Image 파이프라인 설정
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
@@ -59,32 +59,26 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
     torch_dtype=dtype
-)
+).to(device)
 
+# Upscale을 위한 ControlNet 설정
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-)
-
-print("Available attributes in FLUX pipeline:", pipe.__dict__.keys())
+).to(device)
 
-# Upscale 파이프라인 설정 (FLUX 모델의 실제 구조에 맞게 수정)
+# Upscale 파이프라인 설정 (기존 pipe 재사용)
 pipe_upscale = FluxControlNetPipeline(
     vae=pipe.vae,
     text_encoder=pipe.text_encoder,
-    text_encoder_2=pipe.text_encoder_2,
     tokenizer=pipe.tokenizer,
-    tokenizer_2=pipe.tokenizer_2,
     transformer=pipe.transformer,
     scheduler=pipe.scheduler,
     controlnet=controlnet
-)
+).to(device)
+
+
 
-# 추가 속성 설정 (있는 경우에만)
-if hasattr(pipe, 'image_processor'):
-    pipe_upscale.image_processor = pipe.image_processor
 
-# 모든 파이프라인을 device로 이동
-pipe_upscale.to(device)
 
 
 MAX_SEED = 2**32 - 1