Merge pull request #16 from OmarMWarraich/15-save-images-and-write-to-db

15 save images and write to db

Gemfile

@@ -77,3 +77,5 @@ gem 'font-awesome-sass', '~> 5.15.1'
 gem 'ruby-stableDiffusion', path: '/home/owa/Docs/ror/ruby-stablediffusion/'
 
 gem "sidekiq"
+
+gem "ruby-vips", "~> 2.2"

Gemfile.lock

@@ -248,6 +248,9 @@ GEM
       io-console (~> 0.5)
     rexml (3.3.0)
       strscan
+    ruby-vips (2.2.2)
+      ffi (~> 1.12)
+      logger
     rubyzip (2.3.2)
     sassc (2.4.0)
       ffi (~> 1.9)
@@ -329,6 +332,7 @@ DEPENDENCIES
   rails (~> 7.1.3, >= 7.1.3.4)
   redis (>= 4.0.1)
   ruby-stableDiffusion!
+  ruby-vips (~> 2.2)
   sassc-rails (~> 2.1)
   selenium-webdriver
   sidekiq

app/controllers/txt2_imgs_controller.rb

@@ -37,6 +37,7 @@ def create
 
     sd_settings = {
       original_prompt: image_create_params[:prompt],
+      style_template: image_create_params[:style_template],
       prompt: merged_prompt,
       negative_prompt: merged_negative_prompt,
       sampler_name: 'DPM++ 2M Karras',
@@ -55,13 +56,6 @@ def create
       render turbo_stream: turbo_stream.replace(
         'process_starting',
         partial: '/txt2_imgs/process_starting'
-      ) +
-      turbo_stream.replace(
-        'image_maker',
-        partial: '/txt2_imgs/image_maker',
-        locals: {
-          render_result: render_result,
-        }
       )
     end
   end

app/jobs/sd_render_job.rb

@@ -2,8 +2,6 @@ class SdRenderJob < ApplicationJob
   queue_as :default
 
   def perform(render_settings, user_id)
-    original_prompt = render_settings.delete(:original_prompt)
-    style_template = render_settings.delete(:style_template)
 
     config = RStableDiffusionAI::Configuration.new
     config.host = ENV['SD_API_HOST']
@@ -19,9 +17,12 @@ def perform(render_settings, user_id)
     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)
+    ImgProgressJob.perform_later(our_task_id, user_id, render_settings[:original_prompt], render_settings[:style_template])
 
-    result = api_instance.text2imgapi_sdapi_v1_txt2img_post(render_settings)
+    result = api_instance.text2imgapi_sdapi_v1_txt2img_post(render_settings.except(:original_prompt, :style_template))
 
+    GeneratedImage.create_from_sd_render_job(render_settings, user_id, result)
+
+    result
   end
 end

app/models/generated_image.rb

@@ -0,0 +1,60 @@
+# == Schema Information
+#
+# Table name: generated_images
+#
+#  id              :bigint           not null, primary key
+#  info            :json
+#  negative_prompt :string
+#  parameters      :json
+#  prompt          :string
+#  style_template  :string
+#  created_at      :datetime         not null
+#  updated_at      :datetime         not null
+#  user_id         :bigint           not null
+#
+# Indexes
+#
+#  index_generated_images_on_user_id  (user_id)
+#
+# Foreign Keys
+#
+#  fk_rails_...  (user_id => users.id)
+#
+class GeneratedImage < ApplicationRecord
+  include ActionView::RecordIdentifier
+
+  belongs_to :user
+  has_one_attached :sketch
+
+  after_create_commit { broadcast_created }
+
+  def broadcast_created
+    broadcast_update_to(
+      "#{dom_id(user)}_main_image",
+      partial: "generated_images/display_main_image",
+      locals: {
+        scroll_to: true,
+        generated_image: self
+      },
+      target: "image_maker"
+    )
+  end
+
+  def self.create_from_sd_render_job(render_settings, user_id, result)
+    user = User.find(user_id)
+    result.images.each do |src_image|
+      generated_image = user.generated_images.create(
+        prompt: render_settings[:original_prompt],
+        negative_prompt: render_settings[:negative_prompt],
+        style_template: render_settings[:style_template],
+        parameters: result.parameters,
+        info: result.info,
+        sketch: {
+          io: StringIO.new(Base64.decode64(src_image)),
+          content_type: 'image/png',
+          filename: "#{render_settings[:original_prompt]}-#{SecureRandom.uuid}.png".last(256)
+        }
+      )
+    end
+  end
+end

app/models/progress_holder.rb

@@ -4,10 +4,19 @@
 #
 #  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
+#  live_preview_id :integer
+#  user_id         :bigint           not null
+#
+# Indexes
+#
+#  index_progress_holders_on_task_ref  (task_ref) UNIQUE
+#  index_progress_holders_on_user_id   (user_id)
+#
+# Foreign Keys
+#
+#  fk_rails_...  (user_id => users.id)
 #
 class ProgressHolder < ApplicationRecord
   include ActionView::RecordIdentifier

app/models/user.rb

@@ -14,6 +14,8 @@ class User < ApplicationRecord
 
   normalizes :email, with: ->(email) { email.strip.downcase }
 
+  has_many :generated_images, dependent: :destroy
+
   generates_token_for :password_reset, expires_in: 10.minutes do
     password_salt&.last(10)
   end

app/views/generated_images/_display_main_image.html.erb

@@ -0,0 +1 @@
+<%= image_tag generated_image.sketch, alt: "#{generated_image.info}", class: "img-fluid" %>

config/environments/development.rb

@@ -88,3 +88,5 @@
   config.action_mailer.perform_deliveries = true
   config.action_mailer.raise_delivery_errors = true
 end
+
+Rails.application.routes.default_url_options = { host: "localhost", port: 3000 }

db/migrate/20240701214248_create_active_storage_tables.active_storage.rb

@@ -0,0 +1,57 @@
+# This migration comes from active_storage (originally 20170806125915)
+class CreateActiveStorageTables < ActiveRecord::Migration[7.0]
+  def change
+    # Use Active Record's configured type for primary and foreign keys
+    primary_key_type, foreign_key_type = primary_and_foreign_key_types
+
+    create_table :active_storage_blobs, id: primary_key_type do |t|
+      t.string   :key,          null: false
+      t.string   :filename,     null: false
+      t.string   :content_type
+      t.text     :metadata
+      t.string   :service_name, null: false
+      t.bigint   :byte_size,    null: false
+      t.string   :checksum
+
+      if connection.supports_datetime_with_precision?
+        t.datetime :created_at, precision: 6, null: false
+      else
+        t.datetime :created_at, null: false
+      end
+
+      t.index [ :key ], unique: true
+    end
+
+    create_table :active_storage_attachments, id: primary_key_type do |t|
+      t.string     :name,     null: false
+      t.references :record,   null: false, polymorphic: true, index: false, type: foreign_key_type
+      t.references :blob,     null: false, type: foreign_key_type
+
+      if connection.supports_datetime_with_precision?
+        t.datetime :created_at, precision: 6, null: false
+      else
+        t.datetime :created_at, null: false
+      end
+
+      t.index [ :record_type, :record_id, :name, :blob_id ], name: :index_active_storage_attachments_uniqueness, unique: true
+      t.foreign_key :active_storage_blobs, column: :blob_id
+    end
+
+    create_table :active_storage_variant_records, id: primary_key_type do |t|
+      t.belongs_to :blob, null: false, index: false, type: foreign_key_type
+      t.string :variation_digest, null: false
+
+      t.index [ :blob_id, :variation_digest ], name: :index_active_storage_variant_records_uniqueness, unique: true
+      t.foreign_key :active_storage_blobs, column: :blob_id
+    end
+  end
+
+  private
+    def primary_and_foreign_key_types
+      config = Rails.configuration.generators
+      setting = config.options[config.orm][:primary_key_type]
+      primary_key_type = setting || :primary_key
+      foreign_key_type = setting || :bigint
+      [primary_key_type, foreign_key_type]
+    end
+end

db/migrate/20240701221718_create_generated_images.rb

@@ -0,0 +1,14 @@
+class CreateGeneratedImages < ActiveRecord::Migration[7.1]
+  def change
+    create_table :generated_images do |t|
+      t.string :prompt
+      t.string :negative_prompt
+      t.string :style_template
+      t.references :user, null: false, foreign_key: true
+      t.json :info
+      t.json :parameters
+
+      t.timestamps
+    end
+  end
+end

lib/tasks/auto_annotate_models.rake

@@ -0,0 +1,59 @@
+# NOTE: only doing this in development as some production environments (Heroku)
+# NOTE: are sensitive to local FS writes, and besides -- it's just not proper
+# NOTE: to have a dev-mode tool do its thing in production.
+if Rails.env.development?
+  require 'annotate'
+  task :set_annotation_options do
+    # You can override any of these by setting an environment variable of the
+    # same name.
+    Annotate.set_defaults(
+      'active_admin'                => 'false',
+      'additional_file_patterns'    => [],
+      'routes'                      => 'false',
+      'models'                      => 'true',
+      'position_in_routes'          => 'before',
+      'position_in_class'           => 'before',
+      'position_in_test'            => 'before',
+      'position_in_fixture'         => 'before',
+      'position_in_factory'         => 'before',
+      'position_in_serializer'      => 'before',
+      'show_foreign_keys'           => 'true',
+      'show_complete_foreign_keys'  => 'false',
+      'show_indexes'                => 'true',
+      'simple_indexes'              => 'false',
+      'model_dir'                   => 'app/models',
+      'root_dir'                    => '',
+      'include_version'             => 'false',
+      'require'                     => '',
+      'exclude_tests'               => 'false',
+      'exclude_fixtures'            => 'false',
+      'exclude_factories'           => 'false',
+      'exclude_serializers'         => 'false',
+      'exclude_scaffolds'           => 'true',
+      'exclude_controllers'         => 'true',
+      'exclude_helpers'             => 'true',
+      'exclude_sti_subclasses'      => 'false',
+      'ignore_model_sub_dir'        => 'false',
+      'ignore_columns'              => nil,
+      'ignore_routes'               => nil,
+      'ignore_unknown_models'       => 'false',
+      'hide_limit_column_types'     => 'integer,bigint,boolean',
+      'hide_default_column_types'   => 'json,jsonb,hstore',
+      'skip_on_db_migrate'          => 'false',
+      'format_bare'                 => 'true',
+      'format_rdoc'                 => 'false',
+      'format_yard'                 => 'false',
+      'format_markdown'             => 'false',
+      'sort'                        => 'false',
+      'force'                       => 'false',
+      'frozen'                      => 'false',
+      'classified_sort'             => 'true',
+      'trace'                       => 'false',
+      'wrapper_open'                => nil,
+      'wrapper_close'               => nil,
+      'with_comment'                => 'true'
+    )
+  end
+
+  Annotate.load_tasks
+end