Update app.py

app.py

@@ -5,14 +5,18 @@ import gradio as gr
 
 # 模型和 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
+lora_weights_url = "https://huggingface.co/hyder133/chiikawa_stype/resolve/main/tkw1.safetensors"
+
+# 嘗試加載到 GPU，如果不可用則回退到 CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
 # 加載基礎模型
-pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
-pipe = pipe.to("cuda")  # 使用 GPU 加速
+pipe = StableDiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype)
+pipe = pipe.to(device)
 
 # 加載 LoRA 權重
-pipe.load_lora_weights(lora_weights_url)  # 確保權重 URL 正確
+pipe.load_lora_weights(lora_weights_url)
 
 # 測試生成函數
 def generate_image(prompt, width, height, steps, guidance_scale):