Update app.py

app.py

@@ -40,9 +40,14 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
+# FLUX 모델 한 번만 로드
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
+
+# VAE 설정
 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)
+
+# Image2Image 파이프라인 설정
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     base_model,
     vae=good_vae,
@@ -52,12 +57,27 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
     torch_dtype=dtype
-)
+).to(device)
 
 MAX_SEED = 2**32 - 1
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
+# ControlNet 모델과 파이프라인 (필요할 때만 로드)
+controlnet = None
+pipe_controlnet = None
+
+def load_controlnet():
+    global controlnet, pipe_controlnet
+    if controlnet is None:
+        controlnet = FluxControlNetModel.from_pretrained(
+            "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+        ).to(device)
+    if pipe_controlnet is None:
+        pipe_controlnet = FluxControlNetPipeline.from_pretrained(
+            base_model, controlnet=controlnet, torch_dtype=torch.bfloat16
+        ).to(device)
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -488,8 +508,6 @@ css = '''
 footer {visibility: hidden;}
 '''
 
-
-
 huggingface_token = os.getenv("HF_TOKEN")
 
 model_path = snapshot_download(
@@ -497,25 +515,11 @@ model_path = snapshot_download(
     repo_type="model", 
     ignore_patterns=["*.md", "*..gitattributes"],
     local_dir="FLUX.1-dev",
-    token=huggingface_token, # type a new token-id.
+    token=huggingface_token,
 )
 
-# Load pipeline
-controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-).to(device)
-pipe_controlnet = FluxControlNetPipeline.from_pretrained(
-    model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
-)
-pipe_controlnet.to(device)
-
-
-
 MAX_SEED = 1000000
 
-
-
-
 def process_input(input_image, upscale_factor):
     w, h = input_image.size
     w_original, h_original = w, h
@@ -547,12 +551,13 @@ def process_input(input_image, upscale_factor):
 
 MAX_PIXEL_BUDGET = 1024 * 1024
 
-
 @spaces.GPU
 def upscale(input_image, progress=gr.Progress(track_tqdm=True)):
     if input_image is None:
         raise gr.Error("No image to upscale. Please generate an image first.")
     
+    load_controlnet()  # ControlNet 필요 시 로드
+    
     # 입력 이미지 처리
     input_image, w_original, h_original, was_resized = process_input(input_image, 4)
     
@@ -575,8 +580,6 @@ def upscale(input_image, progress=gr.Progress(track_tqdm=True)):
 
     return upscaled_image  # 업스케일된 이미지만 반환
 
-    
-
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as app:
 
     loras_state = gr.State(loras)
@@ -690,8 +693,6 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
         outputs=history_gallery,
     )
 
-
-
     upscale_button.click(
         upscale,
         inputs=[result],
@@ -699,4 +700,4 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as a
     )
 
 app.queue()
-app.launch()
+app.launch()