Update app.py

app.py

@@ -5,7 +5,10 @@ from diffusers import StableDiffusionUpscalePipeline
 
 # 모델 로드
 pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16)
-pipe = pipe.to("cuda")  # GPU 사용
+
+# GPU가 사용 가능한지 확인한 후, 가능하면 GPU로 설정하고, 그렇지 않으면 CPU로 설정
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)  # GPU 또는 CPU 사용
 
 def upscale_image(image, scale_factor):
     if scale_factor == 2:
@@ -17,7 +20,9 @@ def upscale_image(image, scale_factor):
     else:
         raise ValueError("지원하지 않는 배율입니다.")
 
-    upscaled_image = pipe(image=low_res_image).images[0]
+    # 이미지 업스케일링 수행
+    with torch.no_grad():  # 메모리 최적화를 위해 no_grad 사용
+        upscaled_image = pipe(image=low_res_image).images[0]
     return upscaled_image
 
 # Gradio 인터페이스 설정