Update app.py

app.py

@@ -4,26 +4,30 @@ import torch
 from diffusers import StableDiffusionUpscalePipeline
 
 # 모델 로드
-pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
+pipe = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16)
 
 # GPU가 사용 가능한지 확인한 후, 가능하면 GPU로 설정하고, 그렇지 않으면 CPU로 설정
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = pipe.to(device)  # GPU 또는 CPU 사용
+pipe = pipe.to(device)
 
 def upscale_image(image, scale_factor):
+    # 입력 이미지 크기 축소 (예: 가로/세로 각각 절반 크기)
+    reduced_image = image.resize((image.width // 2, image.height // 2), Image.BICUBIC)
+    
     if scale_factor == 2:
         # 2배 업스케일
-        low_res_image = image.resize((image.width // 2, image.height // 2), Image.BICUBIC)
+        low_res_image = reduced_image.resize((reduced_image.width // 2, reduced_image.height // 2), Image.BICUBIC)
     elif scale_factor == 4:
         # 4배 업스케일
-        low_res_image = image.resize((image.width // 4, image.height // 4), Image.BICUBIC)
+        low_res_image = reduced_image.resize((reduced_image.width // 4, reduced_image.height // 4), Image.BICUBIC)
     else:
         raise ValueError("지원하지 않는 배율입니다.")
 
-    # 이미지 업스케일링 수행
+    # 프롬프트 설정 및 업스케일링 수행
     prompt = "high quality, detailed"
     with torch.no_grad():  # 메모리 최적화를 위해 no_grad 사용
         upscaled_image = pipe(prompt=prompt, image=low_res_image).images[0]
+    
     return upscaled_image
 
 # Gradio 인터페이스 설정