Update app.py

app.py

@@ -1,14 +1,18 @@
-from flask import Flask, request, jsonify
-from flask_cors import CORS
-from pymongo import MongoClient
-from bson import ObjectId
 import os
 import oss2
+from pymongo import MongoClient
+from pymongo.server_api import ServerApi
 import bcrypt
-from datetime import datetime, timedelta
-import jwt
-from functools import wraps
-import random
+from datetime import datetime
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import json
+import websocket
+import uuid
+import urllib.request
+import urllib.parse
+import requests
+from bson.objectid import ObjectId
 from dotenv import load_dotenv
 
 # 加载 .env 文件中的环境变量
@@ -17,300 +21,233 @@ load_dotenv()
 app = Flask(__name__)
 CORS(app)
 
-# MongoDB configuration
-# 从环境变量中获取 MongoDB 配置信息
-mongo_uri = os.getenv("MONGO_URI")
-client = MongoClient(mongo_uri)
-db = client['stylespace']
+# ComfyUI 设置
+SERVER_ADDRESS = "127.0.0.1:8188"
+CLIENT_ID = str(uuid.uuid4())
 
-if 'carts' not in db.list_collection_names():
-    db.create_collection('carts')
-    
 # 从环境变量中获取阿里云 OSS 配置信息
 access_key_id = os.getenv("OSS_ACCESS_KEY_ID")
 access_key_secret = os.getenv("OSS_ACCESS_KEY_SECRET")
 bucket_name = os.getenv("OSS_BUCKET_NAME")
 endpoint = os.getenv("OSS_ENDPOINT")
 
+# MongoDB configuration
+# 从环境变量中获取 MongoDB 配置信息
+uri = os.getenv("MONGO_URI")
+client = MongoClient(uri, server_api=ServerApi('1'))
 
 # Create Aliyun OSS bucket object
 bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
 
-# JWT configuration
-app.config['SECRET_KEY'] = '1234'  # Change this to a secure secret key
-
-def token_required(f):
-    @wraps(f)
-    def decorated(*args, **kwargs):
-        token = request.headers.get('Authorization')
-        if not token:
-            return jsonify({'message': 'Token is missing!'}), 401
-        try:
-            token = token.split(" ")[1] if token.startswith("Bearer ") else token
-            data = jwt.decode(token, app.config['SECRET_KEY'], algorithms=["HS256"])
-            current_user = db.users.find_one({'_id': ObjectId(data['user_id'])})
-            if not current_user:
-                raise jwt.InvalidTokenError
-        except jwt.ExpiredSignatureError:
-            return jsonify({'message': 'Token has expired!'}), 401
-        except jwt.InvalidTokenError:
-            return jsonify({'message': 'Invalid token!'}), 401
-        return f(current_user, *args, **kwargs)
-    return decorated
-
-@app.route('/api/auth/register', methods=['POST'])
-def register():
-    data = request.json
-    if db.users.find_one({'username': data['username']}):
-        return jsonify({'message': 'Username already exists'}), 400
-    if db.users.find_one({'email': data['email']}):
-        return jsonify({'message': 'Email already exists'}), 400
-    hashed_password = bcrypt.hashpw(data['password'].encode('utf-8'), bcrypt.gensalt())
-    user = {
-        'username': data['username'],
-        'email': data['email'],
-        'password': hashed_password,
-        'created_at': datetime.utcnow()
+try:
+    client.admin.command('ping')
+    print("Successfully connected to MongoDB!")
+except Exception as e:
+    print(f"Failed to connect to MongoDB: {e}")
+    exit(1)
+
+db = client['ai_image_generator']
+images_collection = db['images']
+
+# ComfyUI workflow (省略具体内容)
+WORKFLOW = {
+    "3": {
+        "inputs": {
+            "seed": 1048756903667323,
+            "steps": 20,
+            "cfg": 8,
+            "sampler_name": "euler",
+            "scheduler": "normal",
+            "denoise": 1,
+            "model": ["4", 0],
+            "positive": ["6", 0],
+            "negative": ["7", 0],
+            "latent_image": ["5", 0]
+        },
+        "class_type": "KSampler"
+    },
+    "4": {
+        "inputs": {
+            "ckpt_name": "sd_xl_base_1.0.safetensors"
+        },
+        "class_type": "CheckpointLoaderSimple"
+    },
+    "5": {
+        "inputs": {
+            "width": 512,
+            "height": 512,
+            "batch_size": 1
+        },
+        "class_type": "EmptyLatentImage"
+    },
+    "6": {
+        "inputs": {
+            "text": "",
+            "clip": ["4", 1]
+        },
+        "class_type": "CLIPTextEncode"
+    },
+    "7": {
+        "inputs": {
+            "text": "text, watermark",
+            "clip": ["4", 1]
+        },
+        "class_type": "CLIPTextEncode"
+    },
+    "8": {
+        "inputs": {
+            "samples": ["3", 0],
+            "vae": ["4", 2]
+        },
+        "class_type": "VAEDecode"
+    },
+    "9": {
+        "inputs": {
+            "filename_prefix": "ComfyUI",
+            "images": ["8", 0]
+        },
+        "class_type": "SaveImage"
     }
-    result = db.users.insert_one(user)
-    return jsonify({'message': 'User created successfully', 'id': str(result.inserted_id)}), 201
-
-@app.route('/api/auth/login', methods=['POST'])
-def login():
-    data = request.json
-    user = db.users.find_one({'username': data['username']})
-    if user and bcrypt.checkpw(data['password'].encode('utf-8'), user['password']):
-        token = jwt.encode(
-            {'user_id': str(user['_id']), 'exp': datetime.utcnow() + timedelta(hours=24)},
-            app.config['SECRET_KEY'],
-            algorithm="HS256"
-        )
-        return jsonify({'token': token, 'user': {'id': str(user['_id']), 'username': user['username'], 'email': user['email']}}), 200
-    return jsonify({'message': 'Invalid username or password'}), 401
-
-@app.route('/api/auth/verify-token', methods=['POST'])
-@token_required
-def verify_token(current_user):
-    return jsonify({'user': {'id': str(current_user['_id']), 'username': current_user['username'], 'email': current_user['email']}}), 200
-
-@app.route('/api/auth/logout', methods=['POST'])
-@token_required
-def logout(current_user):
-    # 由于使用JWT,服务器端不需要做特殊处理
-    return jsonify({'message': 'Successfully logged out'}), 200
+}
+
+def queue_prompt(prompt):
+    p = {"prompt": prompt, "client_id": CLIENT_ID}
+    data = json.dumps(p).encode('utf-8')
+    req = urllib.request.Request(f"http://{SERVER_ADDRESS}/prompt", data=data)
+    return json.loads(urllib.request.urlopen(req).read())
+
+def get_image(filename, subfolder, folder_type):
+    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    url_values = urllib.parse.urlencode(data)
+    with urllib.request.urlopen(f"http://{SERVER_ADDRESS}/view?{url_values}") as response:
+        return response.read()
+
+def get_history(prompt_id):
+    with urllib.request.urlopen(f"http://{SERVER_ADDRESS}/history/{prompt_id}") as response:
+        return json.loads(response.read())
+
+def get_images(ws, prompt):
+    prompt_id = queue_prompt(prompt)['prompt_id']
+    print(f'Prompt ID: {prompt_id}')
+    
+    while True:
+        out = ws.recv()
+        if isinstance(out, str):
+            message = json.loads(out)
+            if message['type'] == 'executing':
+                data = message['data']
+                if data['node'] is None and data['prompt_id'] == prompt_id:
+                    print('Execution completed')
+                    break
+        else:
+            continue  # Ignore binary data (previews)
+
+    history = get_history(prompt_id)[prompt_id]
+    output_images = {}
+
+    for node_id, node_output in history['outputs'].items():
+        if 'images' in node_output:
+            images_output = []
+            for image in node_output['images']:
+                image_data = get_image(image['filename'], image['subfolder'], image['type'])
+                images_output.append(image_data)
+            output_images[node_id] = images_output
+
+    return output_images
+
+def translate_to_english(text):
+    url = os.getenv("DEEP_URI")
+    payload = json.dumps({
+        "text": text,
+        "source_lang": "auto",
+        "target_lang": "EN"
+    })
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    try:
+        response = requests.post(url, headers=headers, data=payload)
+        response.raise_for_status()
+        result = response.json()
+        return result.get('data', text)
+    except requests.RequestException as e:
+        print(f"翻译请求失败: {e}")
+        return text
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+    prompt = request.json['prompt']
+    
+    english_prompt = translate_to_english(prompt)
+    print(f"Original prompt: {prompt}")
+    print(f"Translated prompt: {english_prompt}")
 
-@app.route('/api/products', methods=['GET'])
-def get_products():
-    product_type = request.args.get('type', 'all')
-    limit = int(request.args.get('limit', 20))
+    ws = websocket.create_connection(f"ws://{SERVER_ADDRESS}/ws?clientId={CLIENT_ID}")
 
-    all_products = list(db.products.find())
+    workflow = WORKFLOW.copy()
+    workflow["6"]["inputs"]["text"] = english_prompt
 
-    if product_type == 'hot':
-        selected_products = random.sample(all_products, min(4, len(all_products)))
-    elif product_type == 'new':
-        selected_products = random.sample(all_products, min(4, len(all_products)))
-    else:
-        selected_products = all_products[:limit]
+    images = get_images(ws, workflow)
+    ws.close()
 
-    result = []
-    for product in selected_products:
-        result.append({
-            '_id': str(product['_id']),
-            'name': product.get('name', 'Unknown Product'),
-            'price': float(product.get('price', 0)),
-            'originalPrice': float(product.get('originalPrice', 0)),
-            'brief': product.get('brief', 'No description available'),
-            'image': product.get('images', [None])[0],
-            'description': product.get('description', ''),
-            'images': product.get('images', []),
-            'sizes': product.get('sizes', [])
+    if images:
+        image_data = list(images.values())[0][0]
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+        filename = f"generated_image_{timestamp}.png"
+        oss_path = f"images/{filename}"
+        
+        # Upload to Aliyun OSS
+        bucket.put_object(oss_path, image_data)
+        
+        # Get the public URL
+        image_url = f"https://{bucket_name}.{endpoint}/{oss_path}"
+        
+        # Save to MongoDB
+        image_doc = {
+            "prompt": prompt,
+            "english_prompt": english_prompt,
+            "url": image_url,
+            "filename": filename,
+            "created_at": datetime.utcnow(),
+            "is_public": False
+        }
+        result = images_collection.insert_one(image_doc)
+        
+        return jsonify({
+            "status": "success", 
+            "filename": filename, 
+            "url": image_url, 
+            "id": str(result.inserted_id)
         })
-
-    return jsonify({'data': result}), 200
-
-@app.route('/api/products/<product_id>', methods=['GET'])
-def get_product_details(product_id):
-    product = db.products.find_one({'_id': ObjectId(product_id)})
-    if product:
-        product['_id'] = str(product['_id'])
-        return jsonify({'data': product}), 200
-    return jsonify({'message': 'Product not found'}), 404
-
-@app.route('/api/cart', methods=['GET'])
-@token_required
-def get_cart(current_user):
-    cart = db.carts.find_one({'user_id': str(current_user['_id'])})
-    if cart:
-        cart_items = []
-        for item in cart['items']:
-            product = db.products.find_one({'_id': item['product_id']})
-            if product:
-                cart_items.append({
-                    '_id': str(item['_id']),
-                    'product_id': str(product['_id']),
-                    'name': product['name'],
-                    'price': product['price'],
-                    'image': product['images'][0] if product['images'] else None,
-                    'quantity': item['quantity'],
-                    'size': item['size']
-                })
-        return jsonify({'data': cart_items}), 200
     else:
-        return jsonify({'data': []}), 200
-
-@app.route('/api/cart', methods=['POST'])
-@token_required
-def add_to_cart(current_user):
-    data = request.json
-    cart = db.carts.find_one({'user_id': str(current_user['_id'])})
-    if not cart:
-        cart = {
-            'user_id': str(current_user['_id']),
-            'items': []
-        }
-        result = db.carts.insert_one(cart)
-        cart['_id'] = result.inserted_id
+        return jsonify({"status": "error", "message": "Failed to generate image"})
 
-    item = next((item for item in cart['items'] if str(item['product_id']) == data['productId'] and item['size'] == data['size']), None)
-    if item:
-        item['quantity'] += data['quantity']
+@app.route('/add-to-public-gallery', methods=['POST'])
+def add_to_public_gallery():
+    image_id = request.json['image_id']
+    
+    # Update the image document in MongoDB
+    result = images_collection.update_one(
+        {"_id": ObjectId(image_id)},
+        {"$set": {"is_public": True}}
+    )
+    
+    if result.modified_count > 0:
+        return jsonify({"status": "success", "message": "Image added to public gallery"})
     else:
-        cart['items'].append({
-            '_id': ObjectId(),
-            'product_id': ObjectId(data['productId']),
-            'quantity': data['quantity'],
-            'size': data['size']
-        })
-
-    db.carts.update_one({'_id': cart['_id']}, {'$set': cart})
-
-    return jsonify({'message': 'Item added to cart successfully'}), 200
-
-@app.route('/api/cart/<item_id>', methods=['PUT'])
-@token_required
-def update_cart_item(current_user, item_id):
-    data = request.json
-    cart = db.carts.find_one({'user_id': str(current_user['_id'])})
-    if cart:
-        for item in cart['items']:
-            if str(item['_id']) == item_id:
-                item['quantity'] = data['quantity']
-                break
-        db.carts.update_one({'_id': cart['_id']}, {'$set': cart})
-        
-        updated_item = next((item for item in cart['items'] if str(item['_id']) == item_id), None)
-        if updated_item:
-            product = db.products.find_one({'_id': updated_item['product_id']})
-            if product:
-                return jsonify({
-                    'data': {
-                        '_id': str(updated_item['_id']),
-                        'product_id': str(product['_id']),
-                        'name': product['name'],
-                        'price': product['price'],
-                        'image': product['images'][0] if product['images'] else None,
-                        'quantity': updated_item['quantity'],
-                        'size': updated_item['size']
-                    }
-                }), 200
-        
-        return jsonify({'message': 'Cart item updated successfully'}), 200
-    return jsonify({'message': 'Cart not found'}), 404
-
-@app.route('/api/cart/<item_id>', methods=['DELETE'])
-@token_required
-def remove_from_cart(current_user, item_id):
-    cart = db.carts.find_one({'user_id': str(current_user['_id'])})
-    if cart:
-        cart['items'] = [item for item in cart['items'] if str(item['_id']) != item_id]
-        db.carts.update_one({'_id': cart['_id']}, {'$set': cart})
-        return jsonify({'message': 'Item removed from cart successfully'}), 200
-    return jsonify({'message': 'Cart not found'}), 404
+        return jsonify({"status": "error", "message": "Failed to add image to public gallery"})
 
-@app.route('/api/orders', methods=['POST'])
-@token_required
-def create_order(current_user):
-    data = request.json
-    order = {
-        'user_id': str(current_user['_id']),
-        'orderNumber': generate_order_number(),
-        'items': data['items'],
-        'subtotal': sum(item['price'] * item['quantity'] for item in data['items']),
-        'shippingFee': 10,  # 固定运费,您可以根据需要修改
-        'total': sum(item['price'] * item['quantity'] for item in data['items']) + 10,
-        'status': 'unpaid',
-        'createdAt': datetime.utcnow(),
-        'updatedAt': datetime.utcnow(),
-        'paymentMethod': 'Pending',
-        'shippingAddress': data.get('shippingAddress', {})
-    }
-    result = db.orders.insert_one(order)
-    return jsonify({'message': 'Order created successfully', 'order_id': str(result.inserted_id)}), 201
-
-@app.route('/api/orders/<order_id>', methods=['GET'])
-@token_required
-def get_order_details(current_user, order_id):
-    order = db.orders.find_one({'_id': ObjectId(order_id), 'user_id': str(current_user['_id'])})
-    if order:
-        order['_id'] = str(order['_id'])
-        return jsonify({'data': order}), 200
-    return jsonify({'message': 'Order not found'}), 404
-
-@app.route('/api/orders/<order_id>/status', methods=['PUT'])
-@token_required
-def update_order_status(current_user, order_id):
-    data = request.json
-    new_status = data.get('status')
-    if not new_status:
-        return jsonify({'message': 'Status is required'}), 400
-
-    order = db.orders.find_one({'_id': ObjectId(order_id), 'user_id': str(current_user['_id'])})
-    if not order:
-        return jsonify({'message': 'Order not found'}), 404
-
-    update_data = {
-        'status': new_status,
-        'updatedAt': datetime.utcnow()
-    }
-
-    if new_status == 'unshipped':
-        update_data['paymentMethod'] = data.get('paymentMethod', 'Not specified')
-
-    db.orders.update_one({'_id': ObjectId(order_id)}, {'$set': update_data})
-
-    return jsonify({'message': 'Order status updated successfully'}), 200
-
-@app.route('/api/orders', methods=['GET'])
-@token_required
-def get_user_orders(current_user):
-    status = request.args.get('status')
-    query = {'user_id': str(current_user['_id'])}
-    if status:
-        query['status'] = status
+@app.route('/api/gallery', methods=['GET'])
+def get_gallery_images():
+    # Temporarily return all images, regardless of is_public status
+    all_images = list(images_collection.find().sort("created_at", -1).limit(20))
     
-    orders = list(db.orders.find(query).sort('createdAt', -1))
-    for order in orders:
-        order['_id'] = str(order['_id'])
+    # Convert ObjectId to string for JSON serialization
+    for image in all_images:
+        image['_id'] = str(image['_id'])
     
-    return jsonify({'data': orders}), 200
-
-@app.route('/api/upload', methods=['POST'])
-@token_required
-def upload_image(current_user):
-    if 'file' not in request.files:
-        return jsonify({'message': 'No file part in the request'}), 400
-    file = request.files['file']
-    if file.filename == '':
-        return jsonify({'message': 'No file selected for uploading'}), 400
-    if file:
-        filename = f"uploads/{datetime.now().strftime('%Y%m%d%H%M%S')}_{file.filename}"
-        bucket.put_object(filename, file)
-        url = f"https://{bucket_name}.{endpoint}/{filename}"
-        return jsonify({'message': 'File uploaded successfully', 'url': url}), 200
-
-def generate_order_number():
-    return f"ORD{datetime.utcnow().strftime('%Y%m%d%H%M%S')}{random.randint(1000, 9999)}"
+    print(f"Returning {len(all_images)} images")  # Add this line for debugging
+    return jsonify(all_images)
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860, debug=True)