gradio

app.py

@@ -1,334 +1,69 @@
-from flask import Flask, request, render_template, send_file, jsonify, send_from_directory, session, copy_current_request_context
-from flask_socketio import SocketIO, join_room, leave_room, close_room, rooms, disconnect
-from flask_cors import CORS
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
-import concurrent.futures
-
-import io
+import gradio as gr
 import os
-import argparse
+import io
 from PIL import Image
-import torch
-import gc
-from peft import PeftModel
-
-import queue
-import threading
-import uuid
-import concurrent.futures
-from scripts.process_utils import *
-
-from gevent import pywsgi
-from geventwebsocket.handler import WebSocketHandler
-
-app = Flask(__name__)
-CORS(app)
-socketio = SocketIO(app, cors_allowed_origins="*")
-# クライアントIDとルーム情報を保存するグローバル辞書
-client_rooms = {}
-
-# レート制限の設定
-limiter = Limiter(
-    get_remote_address,
-    app=app,
-    default_limits=["200 per day", "50 per hour"]
-)
-
-# タスクキューの作成とサイズ制限
-MAX_QUEUE_SIZE = 100
-task_queue = queue.Queue(maxsize=MAX_QUEUE_SIZE)
-active_tasks = {}
-task_futures = {}
-
-# ThreadPoolExecutorの作成
-executor = concurrent.futures.ThreadPoolExecutor(max_workers=int(os.environ.get('MAX_WORKERS', 4)))
-
-gpu_lock = threading.Lock()
-
-class Task:
-    def __init__(self, task_id, mode, weight1, weight2, file_data, client_ip, client_id):
-        self.task_id = task_id
-        self.mode = mode
-        self.weight1 = weight1
-        self.weight2 = weight2
-        self.file_data = file_data
-        self.cancel_flag = False
-        self.client_ip = client_ip
-        self.is_processing = False
-        self.client_id = client_id
-
-# キューの状態を通知
-def update_queue_status(message):
-    socketio.emit('queue_update', {'active_tasks': len(active_tasks), 'message': message}, namespace='/demo')
-
-def process_task(task):
-    try:
-        task.is_processing = True
-        # ファイルデータをPIL Imageに変換
-        image = Image.open(io.BytesIO(task.file_data))
-        image = ensure_rgb(image)
-        
-        # キャンセルチェック
-        if task.cancel_flag:
-            return
-
-        # 画像処理ロジックを呼び出す
-        # GPU処理部分
-        with gpu_lock:
-            sotai_image, sketch_image = process_image_as_base64(image, task.mode, task.weight1, task.weight2)
-
-        # キャンセルチェック
-        if task.cancel_flag:
-            return
-        
-        # クライアントIDをリクエストヘッダーから取得（クライアント側で設定する必要があります）
-        client_id = task.client_id
-        if client_id and client_id in client_rooms:
-            room = client_rooms[client_id]
-            # ルームにメッセージをemit
-            socketio.emit('task_complete', {
-                'task_id': task.task_id, 
-                'sotai_image': sotai_image, 
-                'sketch_image': sketch_image
-            }, to=room, namespace='/demo')
-        
-    except Exception as e:
-        print(f"Task error: {str(e)}")
-        if not task.cancel_flag:
-            client_id = task.client_id
-            room = client_rooms[client_id]
-            socketio.emit('task_error', {'task_id': task.task_id, 'error': str(e)}, to=room, namespace='/demo')
-    finally:
-        # タスク数をデクリメント
-        client_ip = task.client_ip
-        tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
-        print(f'Task {task.task_id} completed')
-        task.is_processing = False
-        if task.task_id in active_tasks.keys():
-            del active_tasks[task.task_id]
-        if task.task_id in task_futures.keys():
-            del task_futures[task.task_id]
-
-        update_queue_status('Task completed or cancelled')
-
-def worker():
-    while True:
-        try:
-            task = task_queue.get()
-            if task.task_id in active_tasks.keys():
-                future = executor.submit(process_task, task)
-                task_futures[task.task_id] = future
-        except Exception as e:
-            print(f"Worker error: {str(e)}")
-        finally:
-            # Ensure the task is always removed from the queue
-            task_queue.task_done()
-
-# ワーカースレッドの開始
-threading.Thread(target=worker, daemon=True).start()
-
-# グローバル変数を使用して接続数とタスク数を管理
-connected_clients = 0
-tasks_per_client = {}
-@socketio.on('connect', namespace='/demo')
-def handle_connect(auth):
-    client_id = request.sid
-    room = f"room_{client_id}"  # クライアントごとに一意のルーム名を生成
-    join_room(room)
-    client_rooms[client_id] = room
-    print(f"Client {client_id} connected and joined room {room}")
-    
-    global connected_clients
-    connected_clients += 1
-
-@socketio.on('disconnect' )
-def handle_disconnect():
-    client_id = request.sid
-    if client_id in client_rooms:
-        room = client_rooms[client_id]
-        leave_room(room)
-        del client_rooms[client_id]
-        print(f"Client {client_id} disconnected and removed from room {room}")
-
-    global connected_clients
-    connected_clients -= 1
-    # キャンセル処理：接続が切断された場合、そのクライアントに関連するタスクをキャンセル。ただし、1番目で処理中のタスクはキャンセルしない
-    client_ip = get_remote_address()
-    for task_id, task in active_tasks.items():
-        if task.client_ip == client_ip and not task.is_processing:
-            task.cancel_flag = True
-            if task_id in task_futures:
-                task_futures[task_id].cancel()
-                del task_futures[task_id]
-            del active_tasks[task_id]
-            tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
-
-@app.route('/submit_task', methods=['POST'])
-@limiter.limit("10 per minute")  # 1分間に10回までのリクエストに制限
-def submit_task():
-    if task_queue.full():
-        return jsonify({'error': 'Task queue is full. Please try again later.'}), 503
-
-    # クライアントIPアドレスを取得
-    client_ip = get_remote_address()
-    # 同一IPからの同時タスク数を制限
-    if tasks_per_client.get(client_ip, 0) >= 2:
-        return jsonify({'error': 'Maximum number of concurrent tasks reached'}), 429
-
-    task_id = str(uuid.uuid4())
-    file = request.files['file']
-    mode = request.form.get('mode', 'refine')
-    weight1 = float(request.form.get('weight1', 0.4))
-    weight2 = float(request.form.get('weight2', 0.3))
+import base64
+from scripts.process_utils import initialize, process_image_as_base64
+from scripts.anime import init_model
+from scripts.generate_prompt import load_wd14_tagger_model
+
+# 初期化
+initialize(use_local=False, use_gpu=True)
+init_model(use_local=False)
+load_wd14_tagger_model()
+
+def process_image(input_image, mode, weight1, weight2):
+    # 画像処理ロジック
+    sotai_image, sketch_image = process_image_as_base64(input_image, mode, weight1, weight2)
     
-    # ファイルタイプの制限
-    allowed_extensions = {'png', 'jpg', 'jpeg', 'gif'}
-    if '.' not in file.filename or file.filename.rsplit('.', 1)[1].lower() not in allowed_extensions:
-        return jsonify({'error': 'Invalid file type'}), 415
-
-    # ファイルデータをバイト列として保存
-    file_data = file.read()
+    # Base64文字列をPIL Imageに変換
+    sotai_pil = Image.open(io.BytesIO(base64.b64decode(sotai_image)))
+    sketch_pil = Image.open(io.BytesIO(base64.b64decode(sketch_image)))
     
-    client_id = request.headers.get('X-Client-ID')
-    task = Task(task_id, mode, weight1, weight2, file_data, client_ip, client_id)
-    task_queue.put(task)
-    active_tasks[task_id] = task
+    return sotai_pil, sketch_pil
+
+def gradio_process_image(input_image, mode, weight1, weight2):
+    sotai_image, sketch_image = process_image(input_image, mode, weight1, weight2)
+    return sotai_image, sketch_image
+
+# サンプル画像のパスリスト
+sample_images = [
+    'images/sample1.png',
+    'images/sample2.png',
+    'images/sample4.png',
+    # ... 他のサンプル画像
+]
+
+# Gradio インターフェースの定義
+with gr.Blocks() as demo:
+    gr.Markdown("# Image2Body Test")
     
-    # 同一IPからのタスク数をインクリメント
-    tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) + 1
-
-    update_queue_status('Task submitted') # すべてに通知
-    
-    queue_size = task_queue.qsize()
-    task_order = get_active_task_order(task_id)
-    return jsonify({'task_id': task_id, 'task_order': task_order, 'queue_size': queue_size}) 
-
-@app.route('/cancel_task/<task_id>', methods=['POST'])
-def cancel_task(task_id):
-    # クライアントIPアドレスを取得
-    client_ip = get_remote_address()
-
-    if task_id in active_tasks.keys():
-        task = active_tasks[task_id]
-        # タスクの所有者を確認（IPアドレスで簡易的に判断）
-        if task.client_ip != client_ip:
-            return jsonify({'error': 'Unauthorized to cancel this task'}), 403
-        task.cancel_flag = True
-        if task_id in task_futures.keys():
-            task_futures[task_id].cancel()
-            del task_futures[task_id]
-        del active_tasks[task_id]
-        # タスク数をデクリメント
-        tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
-        update_queue_status('Task cancelled')
-        return jsonify({'message': 'Task cancellation requested'})
-    else:
-        for task in list(task_queue.queue):
-            if task.task_id == task_id and task.client_ip == client_ip:
-                task.cancel_flag = True
-                # タスク数をデクリメント
-                tasks_per_client[client_ip] = tasks_per_client.get(client_ip, 0) - 1
-                return jsonify({'message': 'Task cancellation requested for queued task'})
-        return jsonify({'error': 'Task not found'}), 404
-
-@app.route('/task_status/<task_id>', methods=['GET'])
-def task_status(task_id):
-    try:
-        if task_id in active_tasks.keys():
-            task = active_tasks[task_id]
-            return jsonify({'task_id': task_id, 'is_processing': task.is_processing})
-        else:
-            return jsonify({'task_id': task_id, 'is_processing': False})
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-def get_active_task_order(task_id):
-    try:
-        if task_id not in active_tasks.keys():
-            return 0
-        if active_tasks[task_id].is_processing:
-            return 0
-        processing_task_ids = [tid for tid, task in active_tasks.items() if task.is_processing]
-        non_processing_task_ids = [tid for tid, task in active_tasks.items() if not task.is_processing]
-        if len(processing_task_ids) == 0:
-            task_order = 0
-        else:
-            task_order = non_processing_task_ids.index(task_id) + 1
-        return task_order
-    except Exception as e:
-        print(f"Error getting task order: {str(e)}")
-
-# get_task_orderイベントハンドラー
-@app.route('/get_task_order/<task_id>', methods=['GET'])
-def handle_get_task_order(task_id):
-    if task_id  in active_tasks.keys():
-        return jsonify({'task_order': get_active_task_order(task_id)})
-    else:
-        return jsonify({'task_order': 0})
-
-# Flaskルート
-# ルートパスのGETリクエストに対するハンドラ
-@app.route('/', methods=['GET'])
-def root():
-    return render_template("index.html")
-
-# process_refined のエンドポイント
-@app.route('/process_refined', methods=['POST'])
-def process_refined():
-    file = request.files['file']
-    weight1 = float(request.form.get('weight1', 0.4))
-    weight2 = float(request.form.get('weight2', 0.3))
-    
-    image = ensure_rgb(Image.open(file.stream))
-    sotai_image, sketch_image = process_image_as_base64(image, "refine", weight1, weight2)
-    
-    return jsonify({
-        'sotai_image': sotai_image,
-        'sketch_image': sketch_image
-    })
-
-@app.route('/process_original', methods=['POST'])
-def process_original():
-    file = request.files['file']
-    
-    image = ensure_rgb(Image.open(file.stream))
-    sotai_image, sketch_image = process_image_as_base64(image, "original")
-    
-    return jsonify({
-        'sotai_image': sotai_image,
-        'sketch_image': sketch_image
-    })
-
-@app.route('/process_sketch', methods=['POST'])
-def process_sketch():
-    file = request.files['file']
-    
-    image = ensure_rgb(Image.open(file.stream))
-    sotai_image, sketch_image = process_image_as_base64(image, "sketch")
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image")
+            mode = gr.Radio(["original", "refine"], label="Mode", value="original")
+            with gr.Row():
+                weight1 = gr.Slider(0, 2, value=0.6, step=0.05, label="Weight 1 (Sketch)")
+                weight2 = gr.Slider(0, 1, value=0.05, step=0.025, label="Weight 2 (Body)")
+            process_btn = gr.Button("Process")
+        
+        with gr.Column():
+            sotai_output = gr.Image(type="pil", label="Sotai (Body) Image")
+            sketch_output = gr.Image(type="pil", label="Sketch Image")
     
-    return jsonify({
-        'sotai_image': sotai_image,
-        'sketch_image': sketch_image
-    })
-
-# グローバルエラーハンドラー
-@app.errorhandler(Exception)
-def handle_exception(e):
-    # ログにエラーを記録
-    app.logger.error(f"Unhandled exception: {str(e)}")
-    return jsonify({'error': 'An unexpected error occurred'}), 500
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Server options.')
-    parser.add_argument('--use_local', action='store_true', help='Use local model')
-    parser.add_argument('--use_gpu', action='store_true', help='Set to True to use GPU but if not available, it will use CPU')
-    parser.add_argument('--use_dotenv', action='store_true', help='Use .env file for environment variables')
-
-    args = parser.parse_args()
+    gr.Examples(
+        examples=sample_images,
+        inputs=input_image,
+        outputs=[sotai_output, sketch_output],
+        fn=gradio_process_image,
+        cache_examples=True,
+    )
     
-    initialize(args.use_local, args.use_gpu, args.use_dotenv)
-
-    port = int(os.environ.get('PORT', 7860))
-    server = pywsgi.WSGIServer(('0.0.0.0', port), app, handler_class=WebSocketHandler)
-    server.serve_forever()
+    process_btn.click(
+        fn=gradio_process_image,
+        inputs=[input_image, mode, weight1, weight2],
+        outputs=[sotai_output, sketch_output]
+    )
+
+# Spacesへのデプロイ設定
+demo.launch()

requirements.txt

@@ -2,6 +2,7 @@
 torch==2.2.0
 torchvision==0.17.0
 torchaudio==2.2.0
+transformers
 diffusers==0.27.0 # pth file cannot be loaded in the latest version
 Flask==3.0.3
 Flask-Cors==4.0.0