fix

app.py

@@ -30,6 +30,7 @@ try:
 except Exception as e:
     raise RuntimeError(f"安全模型加载失败: {str(e)}")
 
+
 # 会话管理
 class SessionManager:
     _instances = {}
@@ -50,12 +51,17 @@ class SessionManager:
     def cleanup_sessions(cls):
         with cls._lock:
             now = time.time()
-            expired = [k for k, v in cls._instances.items() if now - v['last_active'] > 3600]
+            expired = [
+                k for k, v in cls._instances.items()
+                if now - v['last_active'] > 3600
+            ]
             for k in expired:
                 del cls._instances[k]
 
+
 # 速率限制
 class RateLimiter:
+
     def __init__(self):
         self.clients = {}
         self.lock = threading.Lock()
@@ -74,9 +80,11 @@ class RateLimiter:
             self.clients[client_id]['count'] += 1
             return True
 
+
 session_manager = SessionManager()
 rate_limiter = RateLimiter()
 
+
 # 工具函数
 def create_error_image(message):
     """生成错误提示图片"""
@@ -90,28 +98,29 @@ def create_error_image(message):
     draw.text((50, 200), text, fill="#ff0000", font=font)
     return img
 
+
 @torch.no_grad()
 def classify_prompt(prompt):
     """安全分类"""
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    inputs = tokenizer(prompt,
+                       return_tensors="pt",
+                       truncation=True,
+                       max_length=512)
     outputs = model(**inputs)
     return torch.argmax(outputs.logits).item()
 
+
 def image_to_base64(file_path):
     """将图片转换为Base64格式"""
     with open(file_path, "rb") as f:
         file_ext = Path(file_path).suffix.lower()[1:]
-        mime_type = f"image/{file_ext}" if file_ext in ["jpeg", "jpg", "png"] else "image/jpeg"
+        mime_type = f"image/{file_ext}" if file_ext in ["jpeg", "jpg", "png"
+                                                        ] else "image/jpeg"
         return f"data:{mime_type};base64,{base64.b64encode(f.read()).decode()}"
 
+
 # 核心生成逻辑
-def generate_image(
-    image_file,
-    prompt,
-    seed,
-    session_id,
-    enable_safety=True
-):
+def generate_image(image_file, prompt, seed, session_id, enable_safety=True):
     try:
         # 安全检查
         if enable_safety:
@@ -123,7 +132,8 @@ def generate_image(
 
         # 速率限制
         if not rate_limiter.check(session_id):
-            error_img = create_error_image("Hourly limit exceeded (20 requests)")
+            error_img = create_error_image(
+                "Hourly limit exceeded (20 requests)")
             yield "❌ 请求过于频繁，请稍后再试", error_img, ""
             return
 
@@ -173,8 +183,7 @@ def generate_image(
             "https://api.wavespeed.ai/api/v2/wavespeed-ai/hidream-e1-full",
             headers=headers,
             json=payload,
-            timeout=30
-        )
+            timeout=30)
         response.raise_for_status()
 
         # 处理响应
@@ -208,83 +217,90 @@ def generate_image(
         error_img = create_error_image(str(e))
         yield f"❌ 生成失败: {str(e)}", error_img, ""
 
+
 # 后台清理任务
 def cleanup_task():
     while True:
         session_manager.cleanup_sessions()
         time.sleep(3600)
 
+
 # 界面构建
-with gr.Blocks(
-    theme=gr.themes.Soft(),
-    css="""
+with gr.Blocks(theme=gr.themes.Soft(),
+               css="""
     .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
     .safe { background: #e8f5e9; border: 1px solid #a5d6a7; }
     .warning { background: #fff3e0; border: 1px solid #ffcc80; }
     .error { background: #ffebee; border: 1px solid #ef9a9a; }
-    """
-) as app:
-    
+    """) as app:
+
     session_id = gr.State(str(uuid.uuid4()))
-    
+
     gr.Markdown("# 🖼️Hidream-E1-Full Live On  Wavespeed Ai")
     gr.Markdown("HiDream-E1 is an image editing model built on HiDream-I1.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            image_file = gr.Image(label="Upload Image", type="filepath", sources=["upload"], interactive=True, image_mode="RGB")
-            prompt = gr.Textbox(label="prompt", placeholder="Please enter an English prompt...", lines=3 )
-            seed = gr.Number(label="seed", value=-1, minimum=-1, maximum=999999, step=1)
+            image_file = gr.Image(label="Upload Image",
+                                  type="filepath",
+                                  sources=["upload"],
+                                  interactive=True,
+                                  image_mode="RGB")
+            prompt = gr.Textbox(
+                label="prompt",
+                placeholder="Please enter an English prompt...",
+                lines=3)
+            seed = gr.Number(label="seed",
+                             value=-1,
+                             minimum=-1,
+                             maximum=999999,
+                             step=1)
             random_btn = gr.Button("random🎲seed", variant="secondary")
-            enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker", value=True, interactive=False)
+            enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",
+                                        value=True,
+                                        interactive=False)
         with gr.Column(scale=1):
             output_image = gr.Image(label="Generated Result")
-            output_url = gr.Textbox(label="image url", interactive=True, visible=False)
+            output_url = gr.Textbox(label="image url",
+                                    interactive=True,
+                                    visible=False)
             status = gr.Textbox(label="Status", elem_classes=["status-box"])
             submit_btn = gr.Button("开始生成", variant="primary")
-    gr.Examples(
-    examples=[
+    gr.Examples(examples=[
         [
-            "Editing Instruction: Convert the image into a Ghibli style. Target Image Description: The woman with long, straight hair and a black top is depicted in a Ghibli style, with vibrant colors and soft, whimsical features, set in a gently lit room.",
-            "https://hidream-ai-hidream-e1-full.hf.space/gradio_api/file=/tmp/gradio/0d28f171c5b089688d5693d7d2bbb02d1daa86ab6c97a8ebf3199dc5f68ab3dc/test_1.png"
+            "Convert the image into Claymation style.",
+            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
         ],
         [
-            "Editing Instruction: Transform the image into a Disney Pixar style. Target Image Description: The person now has stylized, animated features typical of Disney Pixar characters, with exaggerated proportions and vibrant colors, while maintaining the same pose and setting.",
-            "https://hidream-ai-hidream-e1-full.hf.space/gradio_api/file=/tmp/gradio/0d28f171c5b089688d5693d7d2bbb02d1daa86ab6c97a8ebf3199dc5f68ab3dc/test_1.png"
+            "Convert the image into a Ghibli style.",
+            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
         ],
         [
-            "Editing Instruction: Add sunglasses to the girl. Target Image Description: A girl with long brown hair is standing indoors with sunlight streaming through a window, now wearing sunglasses.",
-            "https://hidream-ai-hidream-e1-full.hf.space/gradio_api/file=/tmp/gradio/0d28f171c5b089688d5693d7d2bbb02d1daa86ab6c97a8ebf3199dc5f68ab3dc/test_1.png"
+            "Add sunglasses to the face of the girl.",
+            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl2.png"
         ],
         [
-            'Editing Instruction: Convert the image into an ink sketch style. Target Image Description: An ink sketch of a wooden sign with the text "HiDream E1 Coming Soon" surrounded by ivy, with simplified outlines of mountains and clouds in the background.',
-            "https://hidream-ai-hidream-e1-full.hf.space/gradio_api/file=/tmp/gradio/6cb6be7c9020587a278d8b6d1a7be3e666a96b8f1341361671f9805bd67c2d8b/test_2.jpg"
+            'Convert the image into an ink sketch style.',
+            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
         ],
         [
-            'Editing Instruction: Add a butterfly to the scene. Target Image Description: A wooden sign with the text "HiDream E1 Coming Soon" surrounded by green vines, with a butterfly added to the scene, set against a backdrop of mountains and a clear sky with fluffy clouds.',
-            "https://hidream-ai-hidream-e1-full.hf.space/gradio_api/file=/tmp/gradio/6cb6be7c9020587a278d8b6d1a7be3e666a96b8f1341361671f9805bd67c2d8b/test_2.jpg"
+            'Add a butterfly to the scene.',
+            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_result.png"
         ]
     ],
-    inputs=[prompt, image_file],
-    label="Examples"
-)
+                inputs=[prompt, image_file],
+                label="Examples")
 
-    random_btn.click(
-        fn=lambda: random.randint(0, 999999),
-        outputs=seed
-    )
+    random_btn.click(fn=lambda: random.randint(0, 999999), outputs=seed)
 
     submit_btn.click(
         generate_image,
         inputs=[image_file, prompt, seed, session_id, enable_safety],
-        outputs=[status, output_image, output_url]
-    )
+        outputs=[status, output_image, output_url])
 
 if __name__ == "__main__":
     threading.Thread(target=cleanup_task, daemon=True).start()
-    app.queue(max_size=4).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=False
-    )
+    app.queue(max_size=4).launch(server_name="0.0.0.0",
+                                 server_port=7860,
+                                 show_error=True,
+                                 share=False)