Merge pull request #14 from OmarMWarraich/13-display-image-generation…

…-progress

Display image generation progress

Gemfile

@@ -58,6 +58,8 @@ group :development do
 
   # Speed up commands on slow machines / big apps [https://github.com/rails/spring]
   # gem "spring"
+  gem "annotate"
+  gem "letter_opener"
 end
 
 group :test do
@@ -73,3 +75,5 @@ gem "sassc-rails", "~> 2.1"
 gem 'font-awesome-sass', '~> 5.15.1'
 
 gem 'ruby-stableDiffusion', path: '/home/owa/Docs/ror/ruby-stablediffusion/'
+
+gem "sidekiq"

Gemfile.lock

@@ -84,6 +84,9 @@ GEM
       tzinfo (~> 2.0)
     addressable (2.8.7)
       public_suffix (>= 2.0.2, < 7.0)
+    annotate (3.2.0)
+      activerecord (>= 3.2, < 8.0)
+      rake (>= 10.4, < 14.0)
     autoprefixer-rails (10.4.16.0)
       execjs (~> 2)
     base64 (0.2.0)
@@ -105,6 +108,7 @@ GEM
       rack-test (>= 0.6.3)
       regexp_parser (>= 1.5, < 3.0)
       xpath (~> 3.2)
+    childprocess (5.0.0)
     concurrent-ruby (1.3.3)
     connection_pool (2.4.1)
     crass (1.0.6)
@@ -145,6 +149,11 @@ GEM
       actionview (>= 5.0.0)
       activesupport (>= 5.0.0)
     json (2.7.2)
+    launchy (3.0.1)
+      addressable (~> 2.8)
+      childprocess (~> 5.0)
+    letter_opener (1.10.0)
+      launchy (>= 2.2, < 4)
     logger (1.6.0)
     loofah (2.22.0)
       crass (~> 1.0.2)
@@ -254,6 +263,11 @@ GEM
       rexml (~> 3.2, >= 3.2.5)
       rubyzip (>= 1.2.2, < 3.0)
       websocket (~> 1.0)
+    sidekiq (7.2.4)
+      concurrent-ruby (< 2)
+      connection_pool (>= 2.3.0)
+      rack (>= 2.2.4)
+      redis-client (>= 0.19.0)
     sprockets (4.2.1)
       concurrent-ruby (~> 1.0)
       rack (>= 2.2.4, < 4)
@@ -299,6 +313,7 @@ PLATFORMS
   x86_64-linux
 
 DEPENDENCIES
+  annotate
   bcrypt (~> 3.1.7)
   bootsnap
   bootstrap (~> 5.3)
@@ -308,13 +323,15 @@ DEPENDENCIES
   font-awesome-sass (~> 5.15.1)
   importmap-rails
   jbuilder
+  letter_opener
   pg (~> 1.1)
   puma (>= 5.0)
   rails (~> 7.1.3, >= 7.1.3.4)
   redis (>= 4.0.1)
   ruby-stableDiffusion!
   sassc-rails (~> 2.1)
   selenium-webdriver
+  sidekiq
   sprockets-rails
   stimulus-rails
   turbo-rails

app/controllers/txt2_imgs_controller.rb

@@ -53,10 +53,14 @@ def create
   respond_to do |format|
     format.turbo_stream do
       render turbo_stream: turbo_stream.replace(
+        'process_starting',
+        partial: '/txt2_imgs/process_starting'
+      ) +
+      turbo_stream.replace(
         'image_maker',
-        partial: 'txt2_imgs/image_maker',
+        partial: '/txt2_imgs/image_maker',
         locals: {
-          render_result: render_result
+          render_result: render_result,
         }
       )
     end

app/jobs/img_progress_job.rb

@@ -0,0 +1,35 @@
+class ImgProgressJob < ApplicationJob
+  queue_as :default
+
+  def perform(task_id, user_id, original_prompt, style_template)
+    payload = {
+      "id_task": task_id,
+      "id_live_preview": 1
+    }
+
+    config = RStableDiffusionAI::Configuration.new
+    config.host = ENV['SD_API_HOST']
+    config.scheme = ENV['SD_API_SCHEME']
+    config.debugging = true
+
+    client = RStableDiffusionAI::ApiClient.new(config)
+
+    api_instance = RStableDiffusionAI::DefaultApi.new(client)
+
+    start_time = Time.now
+
+    loop do
+      sleep 5
+      result = api_instance.progressapi_internal_progress_post(payload)
+
+      Rails.logger.debug "Progress: #{result}"
+      break if result.completed == true
+
+      Rails.logger.debug "Time.now - start_time: #{Time.now - start_time}"
+
+      break if Time.now - start_time > 240
+
+      ProgressHolder.new.broadcast_updated(user_id, task_id, original_prompt, style_template, result)
+    end
+  end
+end

app/jobs/sd_render_job.rb

@@ -14,6 +14,13 @@ def perform(render_settings, user_id)
 
     api_instance = RStableDiffusionAI::DefaultApi.new(client)
 
+    our_task_id = "task(#{rand(36**5...36**6).to_s(36)}#{rand(36**5...36**6).to_s(36)}#{rand(36**5...36**6).to_s(36)})"
+
+    render_settings[:task_id] = our_task_id
+    render_settings[:id_task] = our_task_id
+
+    ImgProgressJob.perform_later(our_task_id, user_id, original_prompt, style_template)
+
     result = api_instance.text2imgapi_sdapi_v1_txt2img_post(render_settings)
 
   end

app/models/progress_holder.rb

@@ -0,0 +1,38 @@
+# == Schema Information
+#
+# Table name: progress_holders
+#
+#  id              :bigint           not null, primary key
+#  task_ref        :string
+#  live_preview_id :integer
+#  user_id         :bigint           not null
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#
+class ProgressHolder < ApplicationRecord
+  include ActionView::RecordIdentifier
+  belongs_to :user
+
+  def broadcast_updated(user_id, task_ref, original_prompt, style_template, result)
+    progress_holder = ProgressHolder.find_or_create_by(task_ref:) do |ph|
+      ph.user_id = user_id
+      ph.live_preview_id = result.progress.positive? ? result.id_live_preview : -1
+    end
+
+    user = User.find(user_id)
+
+    # Broadcast the updated progress to the user
+    broadcast_update_to(
+      "#{dom_id(user)}_main_image",
+      partial: "txt2_imgs/loading_progress",
+      locals: {
+        scroll_to: true,
+        original_prompt:,
+        style_template:,
+        result: JSON.parse(result.to_json)
+      },
+      target: "image_maker"
+    )
+
+  end
+end

app/models/user.rb

@@ -1,3 +1,13 @@
+# == Schema Information
+#
+# Table name: users
+#
+#  id              :bigint           not null, primary key
+#  email           :string
+#  password_digest :string
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#
 class User < ApplicationRecord
   has_secure_password
   validates :email, presence: true, uniqueness: true

app/views/txt2_imgs/_loading_progress.html.erb

@@ -0,0 +1,22 @@
+<turbo-frame id="image_maker">
+  <div class="container">
+    <div class="spinner-border mt-3" style="width: 5rem; height: 5rem;" role="status">
+      <span class="visually-hidden">Loading...</span>
+    </div>
+    <br>
+    <div class="m-3 p-3">
+      <% progress = result['progress'].present? ? (result['progress'].to_f * 100) : 0 %>
+      <div class="progress" role="progressbar" aria-label="Basic example" aria-valuenow="<%=progress%>" aria-valuemin="0" aria-valuemax="100">
+        <div class="progress-bar" style="width: <%=progress%>%"></div>
+      </div>
+    </div>
+    <div class="m-3 p-3">
+      <p>ETA: <%= result['eta'].to_f %> seconds</p>
+    </div>
+    <div class="m-3 p-3">
+      <% if result['live_preview'].present? && !result['completed'] != 'true' %>
+        <%= image_tag "#{result['live_preview']}", alt: "#{original_prompt}", class: 'img-fluid thumbnail' %>
+      <% end %>
+    </div>
+  </div>
+</turbo-frame>

app/views/txt2_imgs/_process_starting.html.erb

@@ -0,0 +1,3 @@
+<turbo-frame class="image-maker" action="update">
+  We will make your image in a moment. Please wait...
+</turbo-frame>

app/views/txt2_imgs/index.html.erb

@@ -7,7 +7,7 @@
     </div>
     <div class='col-md-7'>
       <div id='image_holder' class='text-center card h-100 p-2'>
-        <%# turbo_stream_from "#{dom_id(current_user)}_main_image" %>
+        <%= turbo_stream_from "#{dom_id(current_user)}_main_image" %>
         <turbo-frame id="full_image_area">
           <turbo-frame id="image_maker">
             This will be replaced by the image maker

config/application.rb

@@ -23,5 +23,6 @@ class Application < Rails::Application
     #
     # config.time_zone = "Central Time (US & Canada)"
     # config.eager_load_paths << Rails.root.join("extras")
+    config.active_job.queue_adapter = :sidekiq
   end
 end

config/cable.yml

@@ -1,6 +1,6 @@
 development:
   adapter: redis
-  url: redis://localhost:6379/1
+  url: <%= ENV.fetch("REDIS_URL") { "redis://localhost:6379/1" } %>
 
 test:
   adapter: test

db/migrate/20240701172911_create_progress_holders.rb

@@ -0,0 +1,12 @@
+class CreateProgressHolders < ActiveRecord::Migration[7.1]
+  def change
+    create_table :progress_holders do |t|
+      t.string :task_ref
+      t.integer :live_preview_id
+      t.references :user, null: false, foreign_key: true
+
+      t.timestamps
+    end
+    add_index :progress_holders, :task_ref, unique: true
+  end
+end

test/fixtures/progress_holders.yml

@@ -0,0 +1,11 @@
+# Read about fixtures at https://api.rubyonrails.org/classes/ActiveRecord/FixtureSet.html
+
+one:
+  task_ref: MyString
+  live_preview_id: 1
+  user: one
+
+two:
+  task_ref: MyString
+  live_preview_id: 1
+  user: two

test/jobs/img_progress_job_test.rb

@@ -0,0 +1,7 @@
+require "test_helper"
+
+class ImgProgressJobTest < ActiveJob::TestCase
+  # test "the truth" do
+  #   assert true
+  # end
+end

test/models/progress_holder_test.rb

@@ -0,0 +1,18 @@
+# == Schema Information
+#
+# Table name: progress_holders
+#
+#  id              :bigint           not null, primary key
+#  task_ref        :string
+#  live_preview_id :integer
+#  user_id         :bigint           not null
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#
+require "test_helper"
+
+class ProgressHolderTest < ActiveSupport::TestCase
+  # test "the truth" do
+  #   assert true
+  # end
+end

test/models/user_test.rb

@@ -1,3 +1,13 @@
+# == Schema Information
+#
+# Table name: users
+#
+#  id              :bigint           not null, primary key
+#  email           :string
+#  password_digest :string
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#
 require "test_helper"
 
 class UserTest < ActiveSupport::TestCase