Update app.py

app.py

@@ -1,31 +1,51 @@
 from diffusers import StableDiffusionPipeline
+from peft import PeftModel  # PEFT 支持
 import torch
+import gradio as gr
 
-# 模型存儲庫 URL
-model_url = "https://huggingface.co/hyder133/chiikawa_stype/resolve/main/tkw1.safetensors"
+# 模型和 LoRA 權重的 URL
+base_model = "stabilityai/stable-diffusion-xl-base-1.0"  # 基礎模型
+lora_weights_url = "https://huggingface.co/hyder133/chiikawa_stype/resolve/main/tkw1.safetensors"  # 替換為您的 LoRA 權重 URL
 
 # 加載基礎模型
-base_model = "stabilityai/stable-diffusion-xl-base-1.0"  # 基礎模型
 pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # 使用 GPU 加速
 
 # 加載 LoRA 權重
-pipe.load_lora_weights(model_url)  # 指向 LoRA 文件 URL
-pipe = pipe.to("cuda")  # 使用 GPU 加速
+pipe.load_lora_weights(lora_weights_url)  # 確保權重 URL 正確
 
-# 定義生成函數
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
+# 測試生成函數
+def generate_image(prompt, width, height, steps, guidance_scale):
+    # 使用模型生成圖像
+    image = pipe(
+        prompt,
+        width=width,
+        height=height,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale
+    ).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")
+    gr.Markdown("# Chiikawa Style Image Generator")
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt")
+    with gr.Row():
+        width = gr.Slider(256, 1024, value=512, step=64, label="Width")
+        height = gr.Slider(256, 1024, value=512, step=64, label="Height")
+    with gr.Row():
+        steps = gr.Slider(10, 50, value=30, step=5, label="Steps")
+        guidance_scale = gr.Slider(1, 20, value=7.5, step=0.5, label="Guidance Scale")
+    with gr.Row():
+        generate_button = gr.Button("Generate")
+    output = gr.Image(label="Generated Image")
 
-    btn.click(generate_image, inputs=prompt, outputs=output)
+    generate_button.click(
+        generate_image,
+        inputs=[prompt, width, height, steps, guidance_scale],
+        outputs=output
+    )
 
 # 啟動應用
 demo.launch()