Как написать callback для progress bar?

Question

[Владимир Алексеев]

Есть генеративная модель runwayml/stable-diffusion-v1-5, хочу чтобы через websocket передавался на index.html прогресс генерации изображения.

Подскажите, как надо написать progress_callback?

num_inference_steps = 25  # Вы можете настроить это значение

    
    def progress_callback(pipeline, step_index, timestep, callback_kwargs):
        progress = int((step_index + 1) / num_inference_steps * 100)
        socketio.emit('progress', {'progress': progress})
    
    # Генерация изображения
    with torch.no_grad():
        image = pipeline(
            text_prompt,
            num_inference_steps=num_inference_steps,
            callback_on_step_end=progress_callback
        ).images[0]

Comments

[Vindicar]

Ну и в чём проблема? И да, а какой скрипт на index.html это читает и как он это делает?

[Владимир Алексеев]

Vindicar,

<!-- Контейнер для индикатора прогресса -->
        <div
          id="progressContainer"
          class="progress"
          style="display: none; height: 30px; margin-top: 20px"
        >
          <div
            id="progressBar"
            class="progress-bar progress-bar-striped progress-bar-animated"
            role="progressbar"
            style="width: 0%"
          >
            0%
          </div>
        </div>

// Подключаемся к серверу SocketIO
const socket = io();

// Обработчик успешного подключения
socket.on('connect', () => {
    console.log('Connected to Socket.IO server');
});

// Обработчик ошибок
socket.on('connect_error', (error) => {
    console.error('Socket.IO connection error:', error);
});

// Слушаем обновления прогресса
socket.on('progress', function(data) {
    console.log('Received progress:', data.progress);
    const progressBar = document.getElementById('progressBar');
    progressBar.style.width = data.progress + '%';
    progressBar.innerText = data.progress + '%';
});

[Vindicar]

Владимир Алексеев, а как ты сервер socketio выполняешь? Его ведь мало просто создать.

[Владимир Алексеев]

Vindicar, callbeck заработал, надо было return дописать. Теперь socketio.emit('progress', {'progress': progress}) ничего не отправляет на клиент

def progress_callback(pipeline, step_index, timestep, callback_kwargs):
        progress = int(((step_index) / num_inference_steps) * 100)
        socketio.emit('progress', {'progress': progress})
        print(progress)
        return callback_kwargs
    
    # Генерация изображения
    with torch.no_grad():
        image = pipeline(
            text_prompt,
            num_inference_steps=num_inference_steps,
            callback_on_step_end=progress_callback
        ).images[0]