import torch
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(
    "./FLUX.2-9B-bnb-4bit",
    torch_dtype=torch.bfloat16,
    device_map="cuda",                 # REQUIRED
)

#pipe.to("cuda")
#pipe.enable_model_cpu_offload()  # save some VRAM by offloading the model to CPU

from PIL import Image

init_image = Image.open("suji.jpg").convert("RGB")

#prompt = "an very beautiful sexy korean kpop young woman with white bikini is smiling on the waikiki beach. hiqh quality realistic photo."# pixar 3d style"
#prompt = "beautiful woman in the beach holding plate with Circulus "
prompt = "피부가 드러나는 흰색 드레스를 입었다." #하얀색의 섹시한 드레스를 입은 아름다운 한국 여성"
image = pipe(
    prompt=prompt,
    image=init_image,
    height=1024,
    width=1024,
    guidance_scale=1.0,
    num_inference_steps=4,
    generator=torch.Generator(device="cuda").manual_seed(0)
).images[0]
image.save("./output/flux_suji10.png")