Update app.py

app.py

@@ -20,6 +20,9 @@ from gradio_imageslider import ImageSlider
 import numpy as np
 import warnings
 
+
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+
 # 번역 모델 로드
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
@@ -34,21 +37,16 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "black-forest-labs/FLUX.1-dev"
 
-huggingface_token = os.getenv("HUGGINFACE_TOKEN")
-
-model_path = snapshot_download(
-    repo_id="black-forest-labs/FLUX.1-dev", 
-    repo_type="model", 
-    ignore_patterns=["*.md", "*..gitattributes"],
-    local_dir="FLUX.1-dev",
-    token=huggingface_token, # type a new token-id.
-)
+# 공통 FLUX 모델 로드
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, device_map="auto")
 
+# 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)
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
+
+# Image-to-Image 파이프라인 설정
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     base_model,
     vae=good_vae,
@@ -60,14 +58,25 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     torch_dtype=dtype
 )
 
-# Load controlnet for upscaling
+# Upscale을 위한 ControlNet 설정
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
 ).to(device)
-pipe_upscale = FluxControlNetPipeline.from_pretrained(
-    model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
+
+# Upscale 파이프라인 설정 (기존 pipe 재사용)
+pipe_upscale = FluxControlNetPipeline(
+    vae=pipe.vae,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    unet=pipe.unet,
+    scheduler=pipe.scheduler,
+    safety_checker=pipe.safety_checker,
+    feature_extractor=pipe.feature_extractor,
+    controlnet=controlnet
 )
-pipe_upscale.to(device)
+
+
+
 
 MAX_SEED = 2**32 - 1
 MAX_PIXEL_BUDGET = 1024 * 1024