import torch
from diffusers import AutoPipelineForText2Image
from peft import PeftModel, PeftConfig
import gradio as gr # Import Gradio

# 기기 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 기본 모델 로드
print("기본 FLUX 모델 로드 중...")
pipe = AutoPipelineForText2Image.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.float16  # bfloat16 대신 float16 사용
)
pipe.to(device)

# Uncensored LoRA 로드
print("Uncensored LoRA 로드 중...")
pipe.load_lora_weights(
    'Heartsync/Flux-NSFW-uncensored',
    weight_name='lora.safetensors',
    adapter_name="uncensored"
)

# 이미지 생성 함수 정의
def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, width, height, seed):
    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]
    
    # 이미지 저장 (선택 사항, 필요에 따라 활성화/비활성화)
    # image.save("generated_image.png")
    return image

# Gradio 인터페이스 생성
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", value="A woman in a sheer white dress standing on a beach at sunset, backlit so her silhouette is visible through the thin fabric, shot with Canon EOS R5, 85mm f/1.2 lens, golden hour natural lighting, professional composition, hyperrealistic detail, masterpiece quality, 8K resolution."),
        gr.Textbox(label="Negative Prompt", value="text, watermark, signature, cartoon, anime, illustration, painting, drawing, low quality, blurry"),
        gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.0, label="Guidance Scale"),
        gr.Slider(minimum=10, maximum=100, step=1, value=28, label="Number of Inference Steps"),
        gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Width"),
        gr.Slider(minimum=256, maximum=1024, step=64, value=1024, label="Height"),
        gr.Slider(minimum=0, maximum=99999, step=1, value=42, label="Seed")
    ],
    outputs="image",
    title="FLUX.1-dev with Uncensored LoRA",
    description="Generate images using FLUX.1-dev with a loaded Uncensored LoRA model."
)

# 인터페이스 실행 (share=True)
iface.launch(share=True)