Update app.py

app.py

@@ -14,14 +14,34 @@
 
 import dataclasses
 import json
+import base64
+import io
 from pathlib import Path
 
 import gradio as gr
 import torch
 import spaces
+from PIL import Image as PILImage
+from fastapi import FastAPI, Body
+from fastapi.middleware.cors import CORSMiddleware
 
 from uno.flux.pipeline import UNOPipeline
 
+# 创建FastAPI应用
+app = FastAPI()
+
+# 添加CORS中间件允许跨域请求
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 设置全局pipeline变量
+pipeline = None
+
 def get_examples(examples_dir: str = "assets/examples") -> list:
     examples = Path(examples_dir)
     ans = []
@@ -54,6 +74,7 @@ def create_demo(
     device: str = "cuda" if torch.cuda.is_available() else "cpu",
     offload: bool = False,
 ):
+    global pipeline
     pipeline = UNOPipeline(model_type, device, offload, only_lora=True, lora_rank=512)
     pipeline.gradio_generate = spaces.GPU(duratioin=120)(pipeline.gradio_generate)
 
@@ -229,11 +250,117 @@ def create_demo(
                 ],
             )
 
+        # 添加API文档
+        with gr.Accordion("API Documentation", open=False):
+            gr.Markdown("""
+            ### API Usage
+            
+            You can use the following endpoint to generate images programmatically:
+            
+            **Endpoint:** `/api/generate`
+            
+            **Method:** POST
+            
+            **Request Body:**
+            ```json
+            {
+                "prompt": "your text prompt",
+                "image_refs": ["base64_encoded_image1", "base64_encoded_image2", ...],
+                "width": 512,
+                "height": 512,
+                "guidance": 4.0,
+                "num_steps": 25,
+                "seed": -1
+            }
+            ```
+            
+            **Response:**
+            ```json
+            {
+                "image": "base64_encoded_generated_image"
+            }
+            ```
+            
+            **Example JavaScript Usage:**
+            ```javascript
+            async function generateImage() {
+                const response = await fetch('/api/generate', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({
+                        prompt: "handsome woman in the city",
+                        image_refs: [],
+                        width: 512,
+                        height: 512
+                    }),
+                });
+                
+                const data = await response.json();
+                const imgElement = document.getElementById('generatedImage');
+                imgElement.src = `data:image/png;base64,${data.image}`;
+            }
+            ```
+            """)
+
     return demo
 
+# 创建API端点
+@app.post("/api/generate")
+async def generate_image(
+    prompt: str = Body(...),
+    width: int = Body(512),
+    height: int = Body(512),
+    guidance: float = Body(4.0),
+    num_steps: int = Body(25),
+    seed: int = Body(-1),
+    image_refs: list = Body([])
+):
+    global pipeline
+    # 处理参考图像
+    ref_images = []
+    for i in range(min(4, len(image_refs))):
+        if image_refs[i]:
+            try:
+                # 解码base64图像
+                if isinstance(image_refs[i], str) and "base64" in image_refs[i]:
+                    # 移除数据URL前缀
+                    if "," in image_refs[i]:
+                        img_data = image_refs[i].split(",")[1]
+                    else:
+                        img_data = image_refs[i]
+                    
+                    img_data = base64.b64decode(img_data)
+                    ref_img = PILImage.open(io.BytesIO(img_data))
+                    ref_images.append(ref_img)
+                else:
+                    ref_images.append(None)
+            except:
+                ref_images.append(None)
+        else:
+            ref_images.append(None)
+    
+    # 填充至4张图像
+    while len(ref_images) < 4:
+        ref_images.append(None)
+    
+    # 调用模型生成图像
+    result_image, _ = pipeline.gradio_generate(
+        prompt, width, height, guidance, num_steps, seed,
+        ref_images[0], ref_images[1], ref_images[2], ref_images[3]
+    )
+    
+    # 将结果图像编码为base64
+    buffered = io.BytesIO()
+    result_image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    
+    return {"image": img_str}
+
 if __name__ == "__main__":
     from typing import Literal
-
+    import uvicorn
     from transformers import HfArgumentParser
 
     @dataclasses.dataclass
@@ -245,10 +372,17 @@ if __name__ == "__main__":
             metadata={"help": "If True, sequantial offload the models(ae, dit, text encoder) to CPU if not used."}
         )
         port: int = 7860
+        host: str = "0.0.0.0"
 
     parser = HfArgumentParser([AppArgs])
-    args_tuple = parser.parse_args_into_dataclasses() # type: tuple[AppArgs]
+    args_tuple = parser.parse_args_into_dataclasses()  # type: tuple[AppArgs]
     args = args_tuple[0]
-
+    
+    # 创建Gradio demo
     demo = create_demo(args.name, args.device, args.offload)
-    demo.launch(server_port=args.port, ssr_mode=False)
+    
+    # 挂载Gradio接口到FastAPI应用
+    app = gr.mount_gradio_app(app, demo, path="/")
+    
+    # 使用uvicorn启动FastAPI应用
+    uvicorn.run(app, host=args.host, port=args.port)