Update app.py

app.py

@@ -43,11 +43,12 @@ 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).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, low_cpu_mem_usage=True)
+pipe.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)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype, low_cpu_mem_usage=True)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, low_cpu_mem_usage=True)
 
 # Image-to-Image 파이프라인 설정
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
@@ -58,30 +59,26 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     tokenizer=pipe.tokenizer,
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype
-).to(device)
+    torch_dtype=dtype,
+    low_cpu_mem_usage=True
+)
 
 # Upscale을 위한 ControlNet 설정
 controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-).to(device)
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
+)
 
 # Upscale 파이프라인 설정 (기존 pipe 재사용)
 pipe_upscale = FluxControlNetPipeline(
     vae=pipe.vae,
     text_encoder=pipe.text_encoder,
-    text_encoder_2=pipe.text_encoder_2,  # 추가
+    text_encoder_2=pipe.text_encoder_2,
     tokenizer=pipe.tokenizer,
-    tokenizer_2=pipe.tokenizer_2,  # 추가
+    tokenizer_2=pipe.tokenizer_2,
     transformer=pipe.transformer,
     scheduler=pipe.scheduler,
     controlnet=controlnet
-).to(device)
-
-
-
-
-
+)
 
 MAX_SEED = 2**32 - 1
 MAX_PIXEL_BUDGET = 1024 * 1024
@@ -586,16 +583,14 @@ def infer_upscale(
         gr.Info("Upscaling image...")
         # 모든 텐서를 동일한 디바이스로 이동
         pipe_upscale.to(device)
-        control_image = control_image.to(device)
+        control_image = torch.from_numpy(np.array(control_image)).permute(2, 0, 1).float().to(device)
         
         image = pipe_upscale(
             prompt="",
-            control_image=control_image,
+            image=control_image,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
             num_inference_steps=num_inference_steps,
             guidance_scale=3.5,
-            height=control_image.size[1],
-            width=control_image.size[0],
             generator=generator,
         ).images[0]
 
@@ -610,11 +605,11 @@ def infer_upscale(
         return image, seed
     except Exception as e:
         print(f"Error in infer_upscale: {str(e)}")
-        return None, seed
+        return gr.Error(f"Upscaling failed: {str(e)}"), seed
 
 def check_upscale_input(input_image, *args):
     if input_image is None:
-        raise gr.Error("Please provide an input image for upscaling.")
+        return gr.Error("Please provide an input image for upscaling."), *args
     return input_image, *args
     
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as app: