Update app.py

app.py

@@ -1,87 +1,47 @@
 # app.py
-import streamlit as st
-import requests
-import io
-from PIL import Image
-import base64
-import os
+import torch
+import gradio as gr
+from diffusers import DiffusionPipeline
 
-# ----------------------------
-# 页面设置
-# ----------------------------
-st.set_page_config(page_title="StyleDiffusion API Demo", layout="wide")
-st.title("Style Diffusion 推理 Demo (Hugging Face API)")
-st.write("直接调用 Hugging Face 公有模型，无需下载权重")
+print("Loading pipeline...")
 
-# ----------------------------
-# 用户输入
-# ----------------------------
-prompt = st.text_input(
-    "Prompt",
-    "A <new1> reference. New Year image with a rabbit in 2D anime style"
+pipe = DiffusionPipeline.from_pretrained(
+    "GGPENG/StyleDiffusion",
+    torch_dtype=torch.bfloat16,
 )
-steps = st.slider("Steps", 10, 320, 50)
-guidance = st.slider("Guidance Scale", 1.0, 20.0, 7.5)
 
-# ----------------------------
-# Hugging Face API 配置
-# ----------------------------
-API_URL = "https://router.huggingface.co/models/GGPENG/StyleDiffusion"
-API_TOKEN = os.getenv("HF_TOKEN")  # 如果模型是公有的，可以留空
-headers = {"Authorization": f"Bearer {API_TOKEN}"} if API_TOKEN else {}
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ----------------------------
-# 生成函数
-# ----------------------------
-def generate(prompt, steps, guidance):
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "num_inference_steps": steps,
-            "guidance_scale": guidance
-        }
-    }
-
-    response = requests.post(API_URL, headers=headers, json=payload)
+pipe.to("cuda")
 
-    if response.status_code != 200:
-        st.error(f"API请求失败：{response.status_code} {response.text}")
-        return None
+pipe.unet.load_attn_procs(
+"./pytorch_custom_diffusion_weights.bin"
+)
 
-    try:
-        res_json = response.json()
-        if isinstance(res_json, list) and len(res_json) > 0:
-            if "image_base64" in res_json[0]:
-                img_bytes = base64.b64decode(res_json[0]["image_base64"])
-            elif "generated_image" in res_json[0]:
-                img_bytes = base64.b64decode(res_json[0]["generated_image"])
-            else:
-                st.error("API 返回不包含图像字段")
-                return None
-            return Image.open(io.BytesIO(img_bytes))
-        else:
-            st.error("API 返回格式不正确")
-            return None
-    except Exception as e:
-        st.error(f"解析图像失败: {e}")
-        return None
+def generate(prompt, steps, guidance):
 
-# ----------------------------
-# 生成按钮
-# ----------------------------
-if st.button("Generate"):
-    if not prompt.strip():
-        st.warning("请输入 Prompt")
-    else:
-        with st.spinner("生成中，请稍候..."):
-            image = generate(prompt, steps, guidance)
+    image = pipe(
+        prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        eta=1
+    ).images[0]
+
+    return image
+
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Textbox(
+            label="Prompt",
+            value="A <new1> reference. New Year image with a rabbit as the main element"
+        ),
+        gr.Slider(10, 320, value=100, label="Steps"),
+        gr.Slider(1, 18, value=6, label="Guidance"),
+    ],
+    outputs=gr.Image(),
+    title="Fine-tuning style diffusion Demo"
+)
 
-            if image:
-                st.image(image, caption="生成结果", use_column_width=True)
-                buf = io.BytesIO()
-                image.save(buf, format="PNG")
-                st.download_button(
-                    "下载图片",
-                    buf.getvalue(),
-                    "result.png"
-                )
+demo.launch()