from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
import torch
from PIL import Image

def generate_sdxl_with_refiner(prompt: str):
    # 1단계: Base 모델로 초기 이미지 생성
    base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
    base_pipe = StableDiffusionXLPipeline.from_pretrained(
        base_model_id,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True
    ).to("cuda")

    base_image = base_pipe(prompt=prompt, num_inference_steps=30).images[0]

    # 2단계: Refiner 모델로 이미지 개선
    refiner_model_id = "stabilityai/stable-diffusion-xl-refiner-1.0"
    refiner_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        refiner_model_id,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True
    ).to("cuda")

    # Refiner는 PIL 이미지를 입력받아 후처리
    refined_image = refiner_pipe(
        prompt=prompt,
        image=base_image,
        strength=0.3  # 얼마나 많이 보정할지 (0~1)
    ).images[0]

    refined_image.save("output_sdxl_refined.png")
    return refined_image

if __name__ == "__main__":
    prompt = '그녀를 바라보는 한 남자의 야망'
    img = generate_sdxl_with_refiner(prompt)
    img.show()