cli.rs
1 use clap::{Parser, Subcommand}; 2 use funko_diffusion::{Installer, FunkoDiffusionModel, Mode, SessionOptions}; 3 4 #[derive(Parser)] 5 #[command(name = "funko-diffusion")] 6 #[command(about = "Generate Funko Pop style images using diffusion models")] 7 struct Cli { 8 /// Number of threads used within a single op. Defaults to ORT auto-select. 9 #[arg(long)] 10 threads: Option<usize>, 11 12 /// Run VAE encoder/decoder on CPU instead of CUDA. Useful on low-VRAM GPUs. 13 #[cfg(feature = "cuda")] 14 #[arg(long, default_value_t = false)] 15 cpu_vae: bool, 16 17 /// Run UNet on CPU (fp32) instead of CUDA (fp16). Much slower but uses no VRAM. 18 #[cfg(feature = "cuda")] 19 #[arg(long, default_value_t = false)] 20 cpu_unet: bool, 21 22 #[command(subcommand)] 23 command: Commands, 24 } 25 26 #[derive(Subcommand)] 27 enum Commands { 28 #[cfg(feature = "txt2img")] 29 Txt2img { 30 #[arg(short, long)] 31 prompt: String, 32 #[arg(short, long)] 33 negative_prompt: Option<String>, 34 #[arg(short, long, default_value = "output.png")] 35 output: String, 36 #[arg(long, default_value_t = 1024)] 37 width: u32, 38 #[arg(long, default_value_t = 1024)] 39 height: u32, 40 #[arg(long, default_value_t = 40)] 41 steps: usize, 42 #[arg(long, default_value_t = 7.5)] 43 guidance_scale: f32, 44 #[arg(long)] 45 seed: Option<u64>, 46 }, 47 48 #[cfg(feature = "img2img")] 49 Img2img { 50 #[arg(short, long)] 51 prompt: String, 52 #[arg(short, long)] 53 negative_prompt: Option<String>, 54 #[arg(short, long)] 55 input: String, 56 #[arg(short, long, default_value = "output.png")] 57 output: String, 58 #[arg(long)] 59 width: Option<u32>, 60 #[arg(long)] 61 height: Option<u32>, 62 #[arg(long, default_value_t = 0.75)] 63 strength: f32, 64 #[arg(long, default_value_t = 40)] 65 steps: usize, 66 #[arg(long, default_value_t = 7.5)] 67 guidance_scale: f32, 68 #[arg(long)] 69 seed: Option<u64>, 70 }, 71 72 #[cfg(feature = "inpaint")] 73 Inpaint { 74 #[arg(short, long)] 75 prompt: String, 76 #[arg(short, long)] 77 negative_prompt: Option<String>, 78 #[arg(short, long)] 79 input: String, 80 #[arg(short, long)] 81 mask: String, 82 #[arg(short, long, default_value = "output.png")] 83 output: String, 84 #[arg(long)] 85 width: Option<u32>, 86 #[arg(long)] 87 height: Option<u32>, 88 #[arg(long, default_value_t = 0.75)] 89 strength: f32, 90 #[arg(long, default_value_t = 40)] 91 steps: usize, 92 #[arg(long, default_value_t = 7.5)] 93 guidance_scale: f32, 94 #[arg(long)] 95 seed: Option<u64>, 96 }, 97 } 98 99 pub fn run() { 100 let cli = Cli::parse(); 101 let installer = Installer::new("funko-diffusion"); 102 let opts = SessionOptions { 103 intra_op_num_threads: cli.threads, 104 #[cfg(feature = "cuda")] 105 memory_pattern: true, 106 #[cfg(feature = "cuda")] 107 use_cpu_vae: cli.cpu_vae, 108 #[cfg(feature = "cuda")] 109 use_cpu_unet: cli.cpu_unet, 110 }; 111 let mut model = FunkoDiffusionModel::new(&installer, match cli.command { 112 Commands::Txt2img { .. } => Mode::Txt2img, 113 Commands::Img2img { .. } => Mode::Img2img, 114 Commands::Inpaint { .. } => Mode::Inpaint, 115 }, opts).expect("Failed to load model"); 116 117 match cli.command { 118 #[cfg(feature = "txt2img")] 119 Commands::Txt2img { 120 prompt, negative_prompt, output, 121 width, height, steps, guidance_scale, seed, 122 } => { 123 use funko_diffusion::params::TextToImageParams; 124 let image = model.text_to_image(TextToImageParams { 125 prompt, 126 negative_prompt, 127 width, 128 height, 129 steps, 130 guidance_scale, 131 seed, 132 }).expect("Failed to generate image"); 133 image.save(&output).expect("Failed to save image"); 134 println!("Saved to {}", output); 135 } 136 137 #[cfg(feature = "img2img")] 138 Commands::Img2img { 139 prompt, negative_prompt, input, output, 140 width, height, strength, steps, guidance_scale, seed, 141 } => { 142 use funko_diffusion::params::ImageToImageParams; 143 let input_image = image::open(&input).expect("Failed to open input image"); 144 let image = model.image_to_image(ImageToImageParams { 145 prompt, 146 negative_prompt, 147 image: input_image, 148 width, 149 height, 150 strength, 151 steps, 152 guidance_scale, 153 seed, 154 }).expect("Failed to generate image"); 155 image.save(&output).expect("Failed to save image"); 156 println!("Saved to {}", output); 157 } 158 159 #[cfg(feature = "inpaint")] 160 Commands::Inpaint { 161 prompt, negative_prompt, input, mask, output, 162 width, height, strength, steps, guidance_scale, seed, 163 } => { 164 use funko_diffusion::params::InpaintParams; 165 let input_image = image::open(&input).expect("Failed to open input image"); 166 let mask_image = image::open(&mask).expect("Failed to open mask image"); 167 let image = model.inpaint(InpaintParams { 168 prompt, 169 negative_prompt, 170 image: input_image, 171 mask: mask_image, 172 width, 173 height, 174 strength, 175 steps, 176 guidance_scale, 177 seed, 178 }).expect("Failed to generate image"); 179 image.save(&output).expect("Failed to save image"); 180 println!("Saved to {}", output); 181 } 182 } 183 }