import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionUpscalePipeline

# 모델 로드
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)

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 = reduced_image.resize((reduced_image.width // 2, reduced_image.height // 2), Image.BICUBIC)
    elif scale_factor == 4:
        # 4배 업스케일
        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 인터페이스 설정
def main():
    with gr.Blocks() as demo:
        gr.Markdown("# 이미지 업스케일링 웹앱")
        
        # 이미지 업로드
        with gr.Row():
            image_input = gr.Image(type="pil", label="이미지 업로드")

        # 업스케일 배수 선택
        with gr.Row():
            scale_2x_btn = gr.Button("2배 업스케일")
            scale_4x_btn = gr.Button("4배 업스케일")

        # 결과 이미지 출력
        output_image = gr.Image(type="pil", label="업스케일링된 이미지")

        # 버튼 클릭 시 업스케일링 함수 연결
        scale_2x_btn.click(fn=lambda img: upscale_image(img, 2), inputs=image_input, outputs=output_image)
        scale_4x_btn.click(fn=lambda img: upscale_image(img, 4), inputs=image_input, outputs=output_image)

    demo.launch()

if __name__ == "__main__":
    main()