import gradio as gr
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline
import torch
import os

# Hugging Face 토큰 설정
HF_TOKEN = os.getenv("HF_TOKEN")

# 올바른 모델 이름으로 변경 (예시: "lllyasviel/control_v11p_sd15_canny")
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/control_v11p_sd15_canny",
    use_auth_token=HF_TOKEN
)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",  # base model example
    controlnet=controlnet,
    use_auth_token=HF_TOKEN
)

# 이미지 생성 함수
def generate_image(image, prompt, num_inference_steps=50, guidance_scale=7.5):
    if image is None:
        return None
    
    # 이미지에 대한 텍스트 기반 이미지 생성
    result = pipeline(
        prompt=prompt,
        image=image,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).images[0]

    return result

css = """
footer {
    visibility: hidden;
}
"""


# Gradio 인터페이스
demo = gr.Interface(theme="Nymbo/Nymbo_Theme", css=css,
    fn=generate_image,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.Textbox(lines=1, placeholder="Enter a prompt", label="Prompt"),
        gr.Slider(minimum=10, maximum=100, value=50, label="Inference Steps"),
        gr.Slider(minimum=1.0, maximum=20.0, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Stable Diffusion with ControlNet",
    description="Upload an image and enter a prompt to generate a new image with the same face but different context.",
)

if __name__ == "__main__":
    demo.launch()