(mut self)
| 24 | } |
| 25 | |
| 26 | pub async fn run(mut self) -> Result<()> { |
| 27 | if self.ctx.args.api.is_some() { |
| 28 | // run as REST api |
| 29 | api::start(self).await?; |
| 30 | } else { |
| 31 | match self.ctx.args.model_type { |
| 32 | ModelType::TextModel => { |
| 33 | let model = self.model.as_mut().expect("model not found"); |
| 34 | model.add_message(Message::system(self.ctx.args.system_prompt.clone()))?; |
| 35 | model.add_message(Message::user(self.ctx.args.prompt.clone()))?; |
| 36 | |
| 37 | self.generate_text(None, |data| { |
| 38 | if data.is_empty() { |
| 39 | println!(); |
| 40 | } else { |
| 41 | print!("{data}") |
| 42 | } |
| 43 | std::io::stdout().flush().unwrap(); |
| 44 | }) |
| 45 | .await?; |
| 46 | } |
| 47 | ModelType::ImageModel => { |
| 48 | let image_output = self.ctx.args.image_output.clone(); |
| 49 | let mut img_args = self.ctx.args.sd_img_gen_args.clone(); |
| 50 | img_args.set_prompt(&self.ctx.args.prompt); |
| 51 | self.generate_image(img_args, move |images| { |
| 52 | if let Some(image) = images.into_iter().next() { |
| 53 | if let Some(parent) = std::path::Path::new(&image_output).parent() { |
| 54 | if !parent.as_os_str().is_empty() { |
| 55 | std::fs::create_dir_all(parent).ok(); |
| 56 | } |
| 57 | } |
| 58 | image |
| 59 | .save(&image_output) |
| 60 | .expect("Error saving image to disk"); |
| 61 | } |
| 62 | }) |
| 63 | .await?; |
| 64 | } |
| 65 | ModelType::AudioModel => { |
| 66 | let args = AudioGenerationArgs { |
| 67 | input: self.ctx.args.prompt.clone(), |
| 68 | voice_data: None, |
| 69 | voice_path: self.ctx.args.voice_prompt.clone(), |
| 70 | cfg_scale: self.ctx.args.tts_cfg_scale, |
| 71 | max_frames: self.ctx.args.max_audio_frames, |
| 72 | diffusion_steps: self.ctx.args.tts_diffusion_steps, |
| 73 | }; |
| 74 | let output = self.generate_audio(&args).await?; |
| 75 | let wav_bytes = output.to_wav_bytes(); |
| 76 | let output_path = &self.ctx.args.audio_output; |
| 77 | std::fs::write(output_path, &wav_bytes)?; |
| 78 | log::info!( |
| 79 | "Audio saved to {} ({:.1}s, {} samples)", |
| 80 | output_path, |
| 81 | output.samples.len() as f64 / output.sample_rate as f64, |
| 82 | output.samples.len() |
| 83 | ); |
no test coverage detected