Update app.py

app.py

@@ -1,115 +1,168 @@
-import argparse
+import asyncio
+import websockets
+import json
+import logging
+import sys
+import base64
 import os
-import random
-import torch
-import numpy as np
-from diffusers import DiffusionPipeline, AutoencoderKL
-from PIL import Image
-import re
-
-def generate_image(pipe, prompt, seed=42, randomize_seed=True, width=768, height=768, guidance_scale=4.5, num_inference_steps=20):
-    """
-    使用 FLUX.1-Krea-dev 模型生成图像。
-
-    Args:
-        pipe: 配置好的 Diffusers pipeline.
-        prompt (str): 文本提示.
-        seed (int): 随机种子.
-        randomize_seed (bool): 是否随机化种子.
-        width (int): 图像宽度.
-        height (int): 图像高度.
-        guidance_scale (float): 指导比例.
-        num_inference_steps (int): 推理步数.
-
-    Returns:
-        tuple[Image.Image, int]: 返回生成的 PIL 图像和使用的种子.
-    """
-    MAX_SEED = np.iinfo(np.int32).max
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
-    generator = torch.Generator(device=pipe.device).manual_seed(seed)
-    
-    print(f"ℹ️  使用种子: {seed}")
-    print("🚀 开始生成图像...")
-    
-    # 直接调用 pipeline 生成 PIL 图像，内部会自动处理解码
-    image = pipe(
-        prompt=prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-        output_type="pil"
-    ).images[0]
-    
-    return image, seed
+import argparse
 
-def main():
-    """
-    主执行函数，用于解析参数和调用生成逻辑。
-    """
-    # --- 参数解析 ---
-    parser = argparse.ArgumentParser(description="使用 FLUX.1-Krea-dev 模型从文本提示生成图像。")
-    parser.add_argument("--prompt", type=str, required=True, help="用于图像生成的文本提示。")
-    parser.add_argument("--seed", type=int, default=None, help="随机种子。如果未提供，将随机生成。")
-    parser.add_argument("--steps", type=int, default=20, help="推理步数。")
-    parser.add_argument("--width", type=int, default=768, help="图像宽度。")
-    parser.add_argument("--height", type=int, default=768, help="图像高度。")
-    parser.add_argument("--guidance", type=float, default=4.5, help="指导比例 (Guidance Scale)。")
-    args = parser.parse_args()
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-    # --- 模型加载 ---
-    print("⏳ 正在加载模型，请稍候...")
-    dtype = torch.bfloat16
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    # 加载高质量的 VAE 解码器
-    good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", subfolder="vae", torch_dtype=dtype)
-    
-    # 加载主 pipeline，并直接将高质量的 VAE 传入
-    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-Krea-dev", torch_dtype=dtype, vae=good_vae).to(device)
+class FileUploader:
+    def __init__(self, server_url, user_token, space_id):
+        self.server_url = server_url
+        self.user_token = user_token
+        self.space_id = space_id
+        self.websocket = None
+        
+    async def connect_and_upload(self, upload_dir):
+        """连接WebSocket并上传文件"""
+        # 构建WebSocket URL
+        ws_url = self.server_url.replace('http://', 'ws://').replace('https://', 'wss://')
+        uri = f"{ws_url}/ws/upload/{self.space_id}"
+        
+        logger.info(f"正在连接到 {uri}")
+        
+        try:
+            async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
+                self.websocket = websocket
+                logger.info("WebSocket连接成功")
+                
+                # 发送认证信息
+                auth_message = {
+                    "type": "auth",
+                    "token": self.user_token,
+                    "space_id": self.space_id
+                }
+                await websocket.send(json.dumps(auth_message))
+                logger.info("已发送认证信息")
+                
+                # 等待认证响应
+                response = await websocket.recv()
+                auth_result = json.loads(response)
+                
+                if auth_result.get("type") == "auth_success":
+                    logger.info("认证成功，开始上传文件")
+                    await self.upload_directory(upload_dir)
+                else:
+                    logger.error(f"认证失败: {auth_result.get('message', '未知错误')}")
+                    return
+                
+        except Exception as e:
+            logger.error(f"连接失败: {e}")
     
-    if device == "cuda":
-        torch.cuda.empty_cache()
+    async def upload_directory(self, upload_dir):
+        """扫描并上传目录中的所有文件"""
+        if not os.path.exists(upload_dir):
+            logger.error(f"目录不存在: {upload_dir}")
+            return
+        
+        logger.info(f"🔍 开始扫描目录: {upload_dir}")
+        logger.info(f"🔑 Space ID: {self.space_id}")
+        logger.info("-" * 50)
+        
+        # 获取所有文件
+        all_files = []
+        for root, dirs, files in os.walk(upload_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                if os.path.isfile(file_path):
+                    all_files.append(file_path)
+        
+        if not all_files:
+            logger.info("📁 目录中没有找到任何文件")
+            # 发送完成消息
+            await self.websocket.send(json.dumps({
+                "type": "upload_complete",
+                "success_count": 0,
+                "failed_count": 0
+            }))
+            return
+        
+        logger.info(f"📁 找到 {len(all_files)} 个文件，开始上传...")
+        
+        success_count = 0
+        failed_count = 0
+        
+        for file_path in all_files:
+            try:
+                if await self.upload_file(file_path):
+                    success_count += 1
+                else:
+                    failed_count += 1
+                # 稍微延迟一下，避免发送过快
+                await asyncio.sleep(0.1)
+            except Exception as e:
+                logger.error(f"上传文件 {file_path} 时发生异常: {e}")
+                failed_count += 1
+        
+        # 发送上传完成消息
+        await self.websocket.send(json.dumps({
+            "type": "upload_complete",
+            "success_count": success_count,
+            "failed_count": failed_count
+        }))
+        
+        logger.info("-" * 50)
+        logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+        
+        if success_count > 0:
+            logger.info("🎉 文件已成功上传到您的网盘!")
+        
+        return success_count, failed_count
     
-    print(f"✅ 模型加载完成，使用设备: {device}")
+    async def upload_file(self, file_path):
+        """上传单个文件"""
+        if not os.path.exists(file_path):
+            logger.error(f"文件不存在: {file_path}")
+            return False
+        
+        filename = os.path.basename(file_path)
+        logger.info(f"正在上传文件: {filename}")
+        
+        try:
+            with open(file_path, 'rb') as f:
+                file_content = f.read()
+                file_b64 = base64.b64encode(file_content).decode('utf-8')
+                
+                file_message = {
+                    "type": "file_upload",
+                    "filename": filename,
+                    "content": file_b64,
+                    "file_path": os.path.relpath(file_path)
+                }
+                
+                await self.websocket.send(json.dumps(file_message))
+                
+                # 等待上传响应
+                response = await websocket.recv()
+                result = json.loads(response)
+                
+                if result.get("type") == "upload_success":
+                    logger.info(f"✅ 文件 '{filename}' 上传成功!")
+                    return True
+                else:
+                    logger.error(f"❌ 上传失败: {result.get('message', '未知错误')}")
+                    return False
+                    
+        except Exception as e:
+            logger.error(f"上传文件 {file_path} 时出错: {e}")
+            return False
 
-    # --- 图像生成 ---
-    print(f"🎨 开始为提示生成图像: '{args.prompt}'")
+async def main():
+    parser = argparse.ArgumentParser(description="WebSocket文件上传器 - 扫描并上传指定文件夹中的所有文件")
+    parser.add_argument("user_token", help="您的用户令牌(API密钥)")
+    parser.add_argument("space_id", help="Space ID")
+    parser.add_argument("--server", default="ws://127.0.0.1:5001", help="服务器的 WebSocket URL 地址 (默认: ws://127.0.0.1:5001)")
+    parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
     
-    randomize = args.seed is None
-    # 如果用户没有指定种子，则在调用函数时随机化；否则使用用户指定的种子
-    seed_value = args.seed if not randomize else 42 
-
-    generated_image, used_seed = generate_image(
-        pipe=pipe,
-        prompt=args.prompt,
-        seed=seed_value,
-        randomize_seed=randomize,
-        width=args.width,
-        height=args.height,
-        num_inference_steps=args.steps,
-        guidance_scale=args.guidance
-    )
-
-    # --- 保存图像 ---
-    output_dir = "output"
-    os.makedirs(output_dir, exist_ok=True)
-    
-    # 清理提示词以用作安全的文件名
-    safe_prompt = re.sub(r'[^\w\s-]', '', args.prompt).strip()
-    safe_prompt = re.sub(r'[-\s]+', '_', safe_prompt)
+    args = parser.parse_args()
     
-    # 防止文件名过长
-    filename = f"{safe_prompt[:50]}_{used_seed}.png"
-    filepath = os.path.join(output_dir, filename)
-
-    print(f"💾 正在保存图像到: {filepath}")
-    generated_image.save(filepath)
-
-    print("🎉 完成！")
+    uploader = FileUploader(args.server, args.user_token, args.space_id)
+    await uploader.connect_and_upload(args.upload_dir)
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())