Update app.py

app.py

@@ -43,12 +43,11 @@ 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, low_cpu_mem_usage=True)
-pipe.to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
 # LoRA를 위한 설정
-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)
+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,14 +58,13 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
     tokenizer=pipe.tokenizer,
     text_encoder_2=pipe.text_encoder_2,
     tokenizer_2=pipe.tokenizer_2,
-    torch_dtype=dtype,
-    low_cpu_mem_usage=True
-)
+    torch_dtype=dtype
+).to(device)
 
 # Upscale을 위한 ControlNet 설정
 controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
-)
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+).to(device)
 
 # Upscale 파이프라인 설정 (기존 pipe 재사용)
 pipe_upscale = FluxControlNetPipeline(
@@ -78,7 +76,7 @@ pipe_upscale = FluxControlNetPipeline(
     transformer=pipe.transformer,
     scheduler=pipe.scheduler,
     controlnet=controlnet
-)
+).to(device)
 
 MAX_SEED = 2**32 - 1
 MAX_PIXEL_BUDGET = 1024 * 1024
@@ -587,7 +585,7 @@ def infer_upscale(
         
         image = pipe_upscale(
             prompt="",
-            image=control_image,
+            control_image=control_image,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
             num_inference_steps=num_inference_steps,
             guidance_scale=3.5,
@@ -605,7 +603,7 @@ def infer_upscale(
         return image, seed
     except Exception as e:
         print(f"Error in infer_upscale: {str(e)}")
-        return gr.Error(f"Upscaling failed: {str(e)}"), seed
+        return None, seed
 
 def check_upscale_input(input_image, *args):
     if input_image is None: