from diffusers import StableDiffusionPipeline
import torch

# 模型存儲庫 URL
model_url = "https://huggingface.co/hyder133/chiikawa_stype/resolve/main/tkw1.safetensors"

# 加載基礎模型
base_model = "stabilityai/stable-diffusion-xl-base-1.0"  # 基礎模型
pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)

# 加載 LoRA 權重
pipe.load_lora_weights(model_url)  # 指向 LoRA 文件 URL
pipe = pipe.to("cuda")  # 使用 GPU 加速

# 定義生成函數
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio 界面
import gradio as gr
with gr.Blocks() as demo:
    gr.Markdown("## Chiikawa Image Generator")
    prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt")
    output = gr.Image()
    btn = gr.Button("Generate")

    btn.click(generate_image, inputs=prompt, outputs=output)

# 啟動應用
demo.launch()