2023-05-20 17:45:54 +10:00
#frozen_string_literal: true
module DiscourseAi::AiBot::Commands
class ImageCommand < Command
class << self
def name
" image "
end
def desc
2023-08-23 07:49:36 +10:00
" Renders an image from the description (remove all connector words, keep it to 40 words or less). Despite being a text based bot you can generate images! (when user asks to draw, paint or other synonyms try this) "
2023-06-20 08:45:31 +10:00
end
def parameters
[
Parameter . new (
2023-10-27 14:48:12 +11:00
name : " prompts " ,
2023-08-09 16:01:48 +10:00
description :
2023-10-27 14:48:12 +11:00
" The prompts used to generate or create or draw the image (40 words or less, be creative) up to 4 prompts " ,
type : " array " ,
item_type : " string " ,
2023-06-20 08:45:31 +10:00
required : true ,
) ,
2023-10-27 14:48:12 +11:00
Parameter . new (
name : " seeds " ,
description :
" The seed used to generate the image (optional) - can be used to retain image style on amended prompts " ,
type : " array " ,
item_type : " integer " ,
) ,
2023-06-20 08:45:31 +10:00
]
end
2023-05-20 17:45:54 +10:00
end
def result_name
" results "
end
def description_args
2023-06-20 15:44:03 +10:00
{ prompt : @last_prompt }
2023-05-20 17:45:54 +10:00
end
def chain_next_response
2023-08-09 16:01:48 +10:00
false
end
def custom_raw
@custom_raw
2023-05-20 17:45:54 +10:00
end
2023-10-27 14:48:12 +11:00
def process ( prompts : , seeds : nil )
# max 4 prompts
prompts = prompts [ 0 .. 3 ]
seeds = seeds [ 0 .. 3 ] if seeds
@last_prompt = prompts [ 0 ]
2023-08-14 16:30:12 +10:00
show_progress ( localized_description )
results = nil
2023-10-25 11:04:16 +11:00
# this ensures multisite safety since background threads
# generate the images
api_key = SiteSetting . ai_stability_api_key
engine = SiteSetting . ai_stability_engine
api_url = SiteSetting . ai_stability_api_url
2023-10-27 14:48:12 +11:00
threads = [ ]
prompts . each_with_index do | prompt , index |
seed = seeds ? seeds [ index ] : nil
threads << Thread . new ( seed , prompt ) do | inner_seed , inner_prompt |
attempts = 0
begin
DiscourseAi :: Inference :: StabilityGenerator . perform! (
inner_prompt ,
engine : engine ,
api_key : api_key ,
api_url : api_url ,
image_count : 1 ,
seed : inner_seed ,
)
rescue = > e
attempts += 1
retry if attempts < 3
Rails . logger . warn ( " Failed to generate image for prompt #{ prompt } : #{ e } " )
nil
end
2023-08-14 16:30:12 +10:00
end
end
2023-05-20 17:45:54 +10:00
2023-10-27 14:48:12 +11:00
while true
show_progress ( " . " , progress_caret : true )
break if threads . all? { | t | t . join ( 2 ) }
end
results = threads . map ( & :value ) . compact
if ! results . present?
return { prompt : prompt , error : " Something went wrong, could not generate image " }
end
2023-10-25 11:04:16 +11:00
2023-05-20 17:45:54 +10:00
uploads = [ ]
2023-10-27 14:48:12 +11:00
results . each_with_index do | result , index |
result [ :artifacts ] . each do | image |
Tempfile . create ( " v1_txt2img_ #{ index } .png " ) do | file |
file . binmode
file . write ( Base64 . decode64 ( image [ :base64 ] ) )
file . rewind
uploads << {
prompt : prompts [ index ] ,
upload : UploadCreator . new ( file , " image.png " ) . create_for ( bot_user . id ) ,
seed : image [ :seed ] ,
}
end
end
2023-05-20 17:45:54 +10:00
end
2023-08-09 16:01:48 +10:00
@custom_raw = << ~ RAW
2023-06-20 08:45:31 +10:00
[ grid ]
#{
2023-05-20 17:45:54 +10:00
uploads
2023-10-27 14:48:12 +11:00
. map do | item |
" ![ #{ item [ :prompt ] . gsub ( / \ | \ ' \ " / , " " ) } |512x512, 50%]( #{ item [ :upload ] . short_url } ) "
end
2023-05-20 17:45:54 +10:00
. join ( " " )
2023-06-20 08:45:31 +10:00
}
[ / grid]
RAW
2023-10-27 14:48:12 +11:00
{
prompts : uploads . map { | item | { prompt : item [ :prompt ] , seed : item [ :seed ] } } ,
displayed_to_user : true ,
}
2023-05-20 17:45:54 +10:00
end
end
end