feat: add new benchmark Stable Diffusion (#903)

* add scaffold of coco_gen dataset

* test coco_gen

* use smart ptr

* minimal ProcessOutput()

* support float32 output tensor

* Add new benchmark stable_diffusion

* Add dummy_stable_diffusion.tflite

* Add clip_score.cc

* Add generate_tfrecords.py

* Update CaptionRecord

* Add coco_gen.tfrecord

* Compute accuracy

* Format files

* Add notebook to convert CLIP model to TFLite model with resize layer

* Move score prediction from ProcessOutput() to ComputeAccuracy()

* Split clip_score into .h and .cpp files

* Add expected_throughput and expected_accuracy

* Update file paths

* Format files

* refactor TFLite backend to support multiple ML pipelines

* Use custom setting for selecting pipeline

* Update CLIP_Model_to_TFLite.ipynb

* Log record_id for debug purpose

* Refactor benchmark result validity

* Ignore the early-stopping condition for stable_diffusion task

* feat: added stable diffusion pipeline (WIP) (#905)

* feat: added stable diffusion pipeline (WIP)

* - enabled stable diffusion pipeline
- code formatting

* style: code formatting

* style: code formatting

* fix: adjusted inputs and outputs for the stable_diffusion_pipeline

* fix: adjusted sd flow

* fix: added unconditional tokens generation

* chore: adjustments

* Update mode paths for stable_diffusion

* Fix proto/test.cc

* Remove id and image_id from tfrecord file

* Clean up

* Update datasets path

* Delete dummy_stable_diffusion.tflite

* Suggestions for few changes in coco_gen dataset (#908)

* Suggestions for few changes in coco_gen dataset

* Fixing a path

* Update random number seeds

* Format code

* Fix neuron/BUILD file

* Dump raw output to file for COCOGEN (#912)

* Dump output of SD model to file

* Add Python script to convert RGB8 to PNG

* Dump raw output only in accuracy mode

* Disable stable_diffusion task

* Fix Windows build issue

error C2039: '__fs': is not a member of 'std'

* Fix build error on iOS by increasing IPHONEOS_DEPLOYMENT_TARGET to 13.1

error: 'path' is unavailable: introduced in iOS 13.0

* Fix Dart linter error

* Fix build error on iOS by increasing IPHONEOS_DEPLOYMENT_TARGET to 13.1

error: 'path' is unavailable: introduced in iOS 13.0

---------

Co-authored-by: Koan-Sin Tan <[email protected]>
Co-authored-by: RSMNYS <[email protected]>
Co-authored-by: Mohit Mundhra <[email protected]>

Author: 4 people

.bazelrc

@@ -22,6 +22,7 @@ build:ios --copt=-w
 build:windows --copt=/W0
 
 # Build in C++ 17 mode.
+build --cxxopt=-std=c++17
 build:android --cxxopt=-std=c++17
 build:android --host_cxxopt=-std=c++17
 build:ios --cxxopt=-std=c++17

docs/result-spec.md

@@ -104,6 +104,7 @@ If you enable Submission mode, both `performance_run` and `accuracy_run` values
       * `COCO`
       * `ADE20K`
       * `SQUAD`
+      * `COCOGEN`
     * `data_path`: string
     * `groundtruth_path`: string
 * `accuracy_run`: map  

flutter/assets/tasks.pbtxt

@@ -289,3 +289,36 @@ task {
     num_classes: 1000
   }
 }
+
+task {
+  id: "stable_diffusion"
+  name: "Stable Diffusion"
+  min_query_count: 5
+  min_duration: 60
+  max_duration: 600
+  max_throughput: 2000
+  max_accuracy: 1.0
+  scenario: "SingleStream"
+  datasets {
+    type: COCOGEN
+    full {
+      name: "COCO validation set for Stable Diffusion"
+      input_path: "https://github.com/anhappdev/tmp/releases/download/5/coco_gen_test.tfrecord"
+      groundtruth_path: "local:///mlperf_models/stable-diffusion/clip_model_512x512.tflite"
+    }
+    lite {
+      name: "COCO validation set for Stable Diffusion"
+      input_path: "https://github.com/anhappdev/tmp/releases/download/5/coco_gen_full.tfrecord"
+      groundtruth_path: ""
+    }
+    tiny {
+      name: "COCO validation set for Stable Diffusion"
+      input_path: "https://github.com/anhappdev/tmp/releases/download/2/coco_gen.tfrecord"
+      groundtruth_path: "local:///mlperf_models/stable-diffusion/clip_model_512x512.tflite"
+    }
+  }
+  model {
+    id: "stable_diffusion"
+    name: "StableDiffusion"
+  }
+}

flutter/cpp/binary/BUILD

@@ -53,6 +53,7 @@ cc_binary(
         "//flutter/cpp/backends:external",
         "//flutter/cpp/datasets:ade20k",
         "//flutter/cpp/datasets:coco",
+        "//flutter/cpp/datasets:coco_gen",
         "//flutter/cpp/datasets:imagenet",
         "//flutter/cpp/datasets:snu_sr",
         "//flutter/cpp/datasets:squad",

flutter/cpp/binary/main.cc

@@ -24,6 +24,7 @@ limitations under the License.
 #include "flutter/cpp/backends/external.h"
 #include "flutter/cpp/datasets/ade20k.h"
 #include "flutter/cpp/datasets/coco.h"
+#include "flutter/cpp/datasets/coco_gen.h"
 #include "flutter/cpp/datasets/imagenet.h"
 #include "flutter/cpp/datasets/snu_sr.h"
 #include "flutter/cpp/datasets/squad.h"
@@ -64,6 +65,8 @@ DatasetConfig::DatasetType Str2DatasetType(absl::string_view name) {
     return DatasetConfig::ADE20K;
   } else if (absl::EqualsIgnoreCase(name, "SNUSR")) {
     return DatasetConfig::SNUSR;
+  } else if (absl::EqualsIgnoreCase(name, "COCOGEN")) {
+    return DatasetConfig::COCOGEN;
   } else if (absl::EqualsIgnoreCase(name, "DUMMY")) {
     return DatasetConfig::NONE;
   } else {
@@ -83,6 +86,8 @@ DatasetConfig::DatasetType BenchmarkId2DatasetType(absl::string_view name) {
     return DatasetConfig::ADE20K;
   } else if (absl::StartsWith(name, "super_resolution")) {
     return DatasetConfig::SNUSR;
+  } else if (absl::StartsWith(name, "stable_diffusion")) {
+    return DatasetConfig::COCOGEN;
   } else {
     LOG(FATAL) << "Unrecognized benchmark_id: " << name;
     return DatasetConfig::NONE;
@@ -108,7 +113,7 @@ int Main(int argc, char *argv[]) {
           "Benchmark ID. One of image_classification, "
           "image_classification_v2, object_detection, "
           "natural_language_processing, "
-          "image_segmentation_v2, super_resolution, "
+          "image_segmentation_v2, super_resolution, stable_diffusion, "
           "image_classification_offline, image_classification_offline_v2",
           Flag::kPositional)};
   Flags::Parse(&argc, const_cast<const char **>(argv), flag_list);
@@ -362,6 +367,28 @@ int Main(int argc, char *argv[]) {
       flag_list.insert(flag_list.end(), dataset_flags.begin(),
                        dataset_flags.end());
     } break;
+    case DatasetConfig::COCOGEN: {
+      LOG(INFO) << "Using COCO 2014 dataset for Stable Diffusion benchmark";
+      std::string input_tfrecord, input_clip_model = "";
+      std::vector<Flag> dataset_flags{
+          Flag::CreateFlag(
+              "input_tfrecord", &input_tfrecord,
+              "Path to the tfrecord file containing inputs for the model.",
+              Flag::kRequired),
+          Flag::CreateFlag(
+              "input_clip_model", &input_clip_model,
+              "Path to the CLIP model (TFLite) file for score prediction."),
+      };
+
+      if (Flags::Parse(&argc, const_cast<const char **>(argv), dataset_flags) &&
+          backend) {
+        dataset.reset(new CocoGen(backend.get(), input_tfrecord,
+                                  input_clip_model, output_dir));
+      }
+      // Adds to flag_list for showing help.
+      flag_list.insert(flag_list.end(), dataset_flags.begin(),
+                       dataset_flags.end());
+    } break;
     case DatasetConfig::NONE:
     default:
       break;

flutter/cpp/datasets/BUILD

@@ -170,3 +170,37 @@ cc_library(
         "@org_tensorflow//tensorflow/lite/tools/evaluation/stages:image_preprocessing_stage",
     ],
 )
+
+cc_library(
+    name = "coco_gen",
+    srcs = [
+        "coco_gen.cc",
+    ],
+    hdrs = [
+        "coco_gen.h",
+        "utils.h",
+    ],
+    copts = tflite_copts() + select({
+        "//flutter/android/commonlibs:use_asan": [
+            "-fsanitize=address",
+            "-g",
+            "-O1",
+            "-fno-omit-frame-pointer",
+        ],
+        "//conditions:default": [],
+    }),
+    deps = [
+        ":allocator",
+        "//flutter/cpp:mlperf_driver",
+        "//flutter/cpp:utils",
+        "//flutter/cpp/backends:external",
+        "//flutter/cpp/datasets/coco_gen_utils",
+        "//flutter/cpp/datasets/squad_utils",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_protobuf//:protobuf",
+        "@org_tensorflow//tensorflow/lite/tools/evaluation:utils",
+        "@org_tensorflow//tensorflow/lite/tools/evaluation/proto:evaluation_stages_cc_proto",
+        "@org_tensorflow//tensorflow/lite/tools/evaluation/stages:image_preprocessing_stage",
+        "@org_tensorflow//tensorflow/lite/tools/evaluation/stages:object_detection_average_precision_stage",
+    ],
+)

flutter/cpp/datasets/coco_gen.cc

@@ -0,0 +1,156 @@
+/* Copyright 2024 The MLPerf Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "coco_gen.h"
+
+#include <filesystem>
+#include <fstream>
+#include <iomanip>
+
+namespace mlperf {
+namespace mobile {
+
+void dump_output_pixels(const std::vector<uint8_t>& output_pixels,
+                        const std::string& output_filename) {
+  std::ofstream output_file(output_filename, std::ios::binary);
+  if (!output_file) {
+    LOG(ERROR) << "Could not open file for writing: " << output_filename;
+    return;
+  }
+
+  output_file.write(
+      reinterpret_cast<const char*>(output_pixels.data()),
+      static_cast<std::streamsize>(output_pixels.size() * sizeof(uint8_t)));
+
+  if (!output_file) {
+    LOG(ERROR) << "Failed to write data to: " << output_filename;
+  }
+
+  output_file.close();
+  if (!output_file) {
+    LOG(ERROR) << "Could not close the file " << output_filename;
+  }
+  LOG(INFO) << "File saved to: " << output_filename;
+}
+
+CocoGen::CocoGen(Backend* backend, const std::string& input_tfrecord,
+                 const std::string& input_clip_model,
+                 const std::string& output_dir)
+    : Dataset(backend),
+      sample_reader_(input_tfrecord),
+      samples_(sample_reader_.Size()),
+      score_predictor_(input_clip_model) {
+  if (input_format_.size() != 1 || output_format_.size() != 1) {
+    LOG(FATAL) << "Coco_gen only supports 1 input and 1 output";
+    return;
+  }
+
+  isModelFound = score_predictor_.getCanPredict();
+
+  raw_output_dir_ = output_dir + "/cocogen_outputs";
+  std::error_code ec;
+  std::filesystem::create_directories(raw_output_dir_, ec);
+  if (ec) {
+    LOG(ERROR) << "Error: Could not create directory " << raw_output_dir_
+               << ": " << ec.message();
+    return;
+  }
+}
+
+void CocoGen::LoadSamplesToRam(const std::vector<QuerySampleIndex>& samples) {
+  for (QuerySampleIndex sample_idx : samples) {
+    tensorflow::tstring record = sample_reader_.ReadRecord(sample_idx);
+    samples_.at(sample_idx) =
+        std::move(std::make_unique<CaptionRecord>(CaptionRecord(record)));
+  }
+}
+
+void CocoGen::UnloadSamplesFromRam(
+    const std::vector<QuerySampleIndex>& samples) {
+  for (QuerySampleIndex sample_idx : samples) {
+    samples_.at(sample_idx).release();
+  }
+}
+
+#define OUTPUT_SIZE 512 * 512 * 3
+std::vector<uint8_t> CocoGen::ProcessOutput(const int sample_idx,
+                                            const std::vector<void*>& outputs) {
+  if (!isModelFound) return std::vector<uint8_t>();
+  void* output = outputs.at(0);
+  std::vector<uint8_t> output_pixels(OUTPUT_SIZE);
+  if (output_format_[0].type == DataType::Uint8) {
+    uint8_t* temp_data = reinterpret_cast<uint8_t*>(output);
+    std::copy(temp_data, temp_data + output_pixels.size(),
+              output_pixels.begin());
+  } else if (output_format_[0].type == DataType::Float32) {
+    float* temp_data = reinterpret_cast<float*>(output);
+
+    // [-1.0, 1.0] -> [0, 255]
+    for (int i = 0; i < OUTPUT_SIZE; i++) {
+      output_pixels[i] = (uint8_t)((*(temp_data + i) + 1) / 2 * 255);
+    }
+  }
+
+  std::string raw_output_filename =
+      raw_output_dir_ + "/output_" + std::to_string(sample_idx) + ".rgb8";
+  dump_output_pixels(output_pixels, raw_output_filename);
+
+  if (!output_pixels.empty()) {
+    sample_ids_.insert(sample_idx);
+    CaptionRecord* record = samples_.at(sample_idx).get();
+    LOG(INFO) << "caption: " << record->get_caption();
+    caption_map[sample_idx] = record->get_caption();
+    output_pixels_map[sample_idx] = output_pixels;
+    attention_mask_map[sample_idx] = record->get_attention_mask_vector();
+    input_ids_map[sample_idx] = record->get_input_ids_vector();
+    return output_pixels;
+  } else {
+    return std::vector<uint8_t>();
+  }
+}
+
+bool CocoGen::HasAccuracy() { return isModelFound; }
+
+float CocoGen::ComputeAccuracy() {
+  float total_score = 0.0f;
+  float total_samples = static_cast<float>(sample_ids_.size());
+  for (int sample_idx : sample_ids_) {
+    std::string caption = caption_map[sample_idx];
+    std::vector<int32_t> input_ids = input_ids_map[sample_idx];
+    std::vector<int32_t> attention_mask = attention_mask_map[sample_idx];
+    std::vector<uint8_t> output_pixels = output_pixels_map[sample_idx];
+    std::vector<float> pixel_values(OUTPUT_SIZE);
+    for (int i = 0; i < OUTPUT_SIZE; i++) {
+      pixel_values[i] = static_cast<float>((output_pixels[i] / 128.0) - 1.0);
+    }
+    float score =
+        score_predictor_.predict(attention_mask, input_ids, pixel_values);
+    LOG(INFO) << "sample_idx: " << sample_idx << " caption: " << caption
+              << " score: " << score;
+    total_score += score;
+  }
+  float avg_score = total_score / total_samples;
+  return avg_score;
+}
+
+std::string CocoGen::ComputeAccuracyString() {
+  float result = ComputeAccuracy();
+  if (result < 0.0f) {
+    return {"N/A"};
+  }
+  std::stringstream stream;
+  stream << std::fixed << std::setprecision(4) << result;
+  return stream.str();
+}
+
+}  // namespace mobile
+}  // namespace mlperf