Skip to content

Commit fe25b69

Browse files
BH3GEIBH3GEI
authored
Merge pull request #60 from mofa-org/feat/gpt-image-all-pipelines
feat: gpt-image-2 support for all image pipelines + plugin file delivery
2 parents cf37ca4 + b4fceae commit fe25b69

4 files changed

Lines changed: 211 additions & 71 deletions

File tree

‎mofa-cli/src/main.rs‎

Lines changed: 66 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,9 @@ enum Commands {
9999
/// API mode: rt (realtime, default) or batch (50% cheaper, async)
100100
#[arg(long, value_enum, default_value = "rt")]
101101
api: ApiMode,
102+
/// Image generation model (e.g. gpt-image-2)
103+
#[arg(long)]
104+
gen_model: Option<String>,
102105
/// Input JSON file (or stdin)
103106
#[arg(long, short)]
104107
input: Option<PathBuf>,
@@ -132,6 +135,9 @@ enum Commands {
132135
/// API mode: rt (realtime, default) or batch (50% cheaper, async)
133136
#[arg(long, value_enum, default_value = "rt")]
134137
api: ApiMode,
138+
/// Image generation model (e.g. gpt-image-2)
139+
#[arg(long)]
140+
gen_model: Option<String>,
135141
/// Input JSON file (or stdin)
136142
#[arg(long, short)]
137143
input: Option<PathBuf>,
@@ -165,6 +171,9 @@ enum Commands {
165171
/// API mode: rt (realtime, default) or batch (50% cheaper, async)
166172
#[arg(long, value_enum, default_value = "rt")]
167173
api: ApiMode,
174+
/// Image generation model (e.g. gpt-image-2)
175+
#[arg(long)]
176+
gen_model: Option<String>,
168177
/// Input JSON file (or stdin)
169178
#[arg(long, short)]
170179
input: Option<PathBuf>,
@@ -224,6 +233,9 @@ enum Commands {
224233
/// API mode: rt (realtime, default) or batch (50% cheaper, async)
225234
#[arg(long, value_enum, default_value = "rt")]
226235
api: ApiMode,
236+
/// Image generation model (e.g. gpt-image-2)
237+
#[arg(long)]
238+
gen_model: Option<String>,
227239
/// Input JSON file (or stdin)
228240
#[arg(long, short)]
229241
input: Option<PathBuf>,
@@ -277,8 +289,19 @@ fn find_styles_dir(mofa_root: &std::path::Path, skill_name: &str) -> PathBuf {
277289
mofa_root.join("mofa").join("styles")
278290
}
279291

292+
struct PluginOutput {
293+
text: String,
294+
files: Vec<String>,
295+
}
296+
297+
impl From<String> for PluginOutput {
298+
fn from(text: String) -> Self {
299+
Self { text, files: vec![] }
300+
}
301+
}
302+
280303
/// Plugin protocol mode: called as `./main <tool_name>` with JSON on stdin.
281-
/// Returns `{"output": "...", "success": true/false, "summary": ...}` on stdout.
304+
/// Returns `{"output": "...", "success": true/false, "files_to_send": [...]}` on stdout.
282305
fn run_plugin(tool_name: &str, cancel: &std::sync::atomic::AtomicBool) -> Result<()> {
283306
use crate::protocol_v2::{check_cancel, emit_v2_progress};
284307

@@ -322,27 +345,30 @@ fn run_plugin(tool_name: &str, cancel: &std::sync::atomic::AtomicBool) -> Result
322345
// emit a structured summary return `serde_json::Value::Null` —
323346
// the host treats missing summary as "no structured info, use the
324347
// existing output text".
325-
let result: Result<(String, serde_json::Value)> = match tool_name {
326-
"mofa_slides" => plugin_slides(&args, &mofa_root, &cfg, cancel),
327-
"mofa_cards" => plugin_cards(&args, &mofa_root, &cfg).map(|s| (s, serde_json::Value::Null)),
328-
"mofa_comic" => plugin_comic(&args, &mofa_root, &cfg).map(|s| (s, serde_json::Value::Null)),
348+
let result: Result<(PluginOutput, serde_json::Value)> = match tool_name {
349+
"mofa_slides" => plugin_slides(&args, &mofa_root, &cfg, cancel).map(|(s, v)| (s.into(), v)),
350+
"mofa_cards" => plugin_cards(&args, &mofa_root, &cfg).map(|o| (o, serde_json::Value::Null)),
351+
"mofa_comic" => plugin_comic(&args, &mofa_root, &cfg).map(|s| (s.into(), serde_json::Value::Null)),
329352
"mofa_infographic" => {
330-
plugin_infographic(&args, &mofa_root, &cfg).map(|s| (s, serde_json::Value::Null))
353+
plugin_infographic(&args, &mofa_root, &cfg).map(|s| (s.into(), serde_json::Value::Null))
331354
}
332-
"mofa_video" => plugin_video(&args, &mofa_root, &cfg).map(|s| (s, serde_json::Value::Null)),
355+
"mofa_video" => plugin_video(&args, &mofa_root, &cfg).map(|s| (s.into(), serde_json::Value::Null)),
333356
_ => Err(eyre::eyre!("unknown tool: {tool_name}")),
334357
};
335358

336359
match result {
337360
Ok((output, summary)) => {
338361
emit_v2_progress("complete", &format!("{tool_name} complete"), Some(1.0));
339362
let mut payload = serde_json::json!({
340-
"output": output,
363+
"output": output.text,
341364
"success": true,
342365
});
343366
if !summary.is_null() {
344367
payload["summary"] = summary;
345368
}
369+
if !output.files.is_empty() {
370+
payload["files_to_send"] = serde_json::json!(output.files);
371+
}
346372
println!("{payload}");
347373
}
348374
Err(e) => {
@@ -586,7 +612,7 @@ fn plugin_cards(
586612
args: &serde_json::Value,
587613
mofa_root: &std::path::Path,
588614
cfg: &config::MofaConfig,
589-
) -> Result<String> {
615+
) -> Result<PluginOutput> {
590616
let style_name = args
591617
.get("style")
592618
.and_then(|v| v.as_str())
@@ -615,23 +641,37 @@ fn plugin_cards(
615641
std::fs::create_dir_all(&card_dir).ok();
616642

617643
let batch = args.get("api").and_then(|v| v.as_str()).unwrap_or("rt") == "batch";
618-
pipeline::cards::run(
644+
let gen_model = args.get("gen_model").and_then(|v| v.as_str());
645+
let results = pipeline::cards::run(
619646
&card_dir,
620647
&cards,
621648
&loaded_style,
622649
cfg,
623650
concurrency,
624651
aspect,
625652
image_size,
626-
None,
653+
gen_model,
627654
batch,
628655
)?;
629656

630-
Ok(format!(
631-
"Generated {} card(s) in {}",
632-
cards.len(),
633-
card_dir.display()
634-
))
657+
let cwd = std::env::current_dir().unwrap_or_default();
658+
let files: Vec<String> = results
659+
.iter()
660+
.filter_map(|p| {
661+
p.as_ref().map(|p| {
662+
if p.is_absolute() {
663+
p.display().to_string()
664+
} else {
665+
cwd.join(p).display().to_string()
666+
}
667+
})
668+
})
669+
.collect();
670+
671+
Ok(PluginOutput {
672+
text: format!("Generated {} card(s) in {}", cards.len(), card_dir.display()),
673+
files,
674+
})
635675
}
636676

637677
fn plugin_comic(
@@ -835,6 +875,7 @@ fn plugin_video(
835875
std::fs::create_dir_all(&card_dir).ok();
836876

837877
let batch = args.get("api").and_then(|v| v.as_str()).unwrap_or("rt") == "batch";
878+
let gen_model = args.get("gen_model").and_then(|v| v.as_str());
838879
pipeline::video::run(
839880
&card_dir,
840881
&cards,
@@ -851,6 +892,7 @@ fn plugin_video(
851892
music_volume,
852893
music_fade_in,
853894
batch,
895+
gen_model,
854896
)?;
855897

856898
Ok(format!(
@@ -931,6 +973,7 @@ fn main() -> Result<()> {
931973
concurrency,
932974
image_size,
933975
api,
976+
gen_model,
934977
input,
935978
} => {
936979
let styles_dir = find_styles_dir(&mofa_root, "cards");
@@ -948,7 +991,7 @@ fn main() -> Result<()> {
948991
concurrency,
949992
aspect.as_deref(),
950993
image_size.as_deref(),
951-
None,
994+
gen_model.as_deref(),
952995
matches!(api, ApiMode::Batch),
953996
)?;
954997
}
@@ -962,6 +1005,7 @@ fn main() -> Result<()> {
9621005
refine,
9631006
gutter,
9641007
api,
1008+
gen_model,
9651009
input,
9661010
} => {
9671011
let styles_dir = find_styles_dir(&mofa_root, "comic");
@@ -988,7 +1032,7 @@ fn main() -> Result<()> {
9881032
image_size.as_deref(),
9891033
refine,
9901034
gutter,
991-
None,
1035+
gen_model.as_deref(),
9921036
matches!(api, ApiMode::Batch),
9931037
)?;
9941038
}
@@ -1002,6 +1046,7 @@ fn main() -> Result<()> {
10021046
refine,
10031047
gutter,
10041048
api,
1049+
gen_model,
10051050
input,
10061051
} => {
10071052
let styles_dir = find_styles_dir(&mofa_root, "infographic");
@@ -1028,7 +1073,7 @@ fn main() -> Result<()> {
10281073
aspect.as_deref(),
10291074
refine,
10301075
gutter,
1031-
None,
1076+
gen_model.as_deref(),
10321077
matches!(api, ApiMode::Batch),
10331078
)?;
10341079
}
@@ -1046,6 +1091,7 @@ fn main() -> Result<()> {
10461091
music_volume,
10471092
music_fade_in,
10481093
api,
1094+
gen_model,
10491095
input,
10501096
} => {
10511097
let styles_dir = find_styles_dir(&mofa_root, "video");
@@ -1078,6 +1124,7 @@ fn main() -> Result<()> {
10781124
music_volume,
10791125
music_fade_in,
10801126
matches!(api, ApiMode::Batch),
1127+
gen_model.as_deref(),
10811128
)?;
10821129
}
10831130
Commands::PptxUnpack { input, output_dir } => {

‎mofa-cli/src/pipeline/comic.rs‎

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::config::MofaConfig;
44
use crate::dashscope::DashscopeClient;
55
use crate::gemini::{BatchImageRequest, GeminiClient};
66
use crate::image_util;
7+
use crate::openai::OpenAIImageClient;
78
use crate::style::Style;
89
use eyre::Result;
910
use serde::Deserialize;
@@ -20,7 +21,8 @@ pub struct PanelInput {
2021
/// Generate panels using synchronous parallel calls.
2122
#[allow(clippy::too_many_arguments)]
2223
fn gen_panels_sync(
23-
gemini: &GeminiClient,
24+
gemini: &Option<GeminiClient>,
25+
openai: &Option<OpenAIImageClient>,
2426
out_dir: &Path,
2527
panels: &[PanelInput],
2628
style: &Style,
@@ -56,15 +58,35 @@ fn gen_panels_sync(
5658
let padded = format!("{:02}", idx + 1);
5759
let out_path = out_dir.join(format!("panel-{padded}.png"));
5860

59-
if let Ok(Some(p)) = gemini.gen_image(
60-
&full_prompt,
61-
&out_path,
62-
image_size,
63-
Some(panel_aspect),
64-
&[],
65-
Some(model),
66-
Some(&format!("Panel {}", idx + 1)),
67-
) {
61+
let result = if model.starts_with("gpt-image") {
62+
openai.as_ref().and_then(|oa| {
63+
oa.gen_image(
64+
&full_prompt,
65+
&out_path,
66+
image_size,
67+
Some(panel_aspect),
68+
Some(model),
69+
Some(&format!("Panel {}", idx + 1)),
70+
)
71+
.ok()
72+
.flatten()
73+
})
74+
} else {
75+
gemini.as_ref().and_then(|gem| {
76+
gem.gen_image(
77+
&full_prompt,
78+
&out_path,
79+
image_size,
80+
Some(panel_aspect),
81+
&[],
82+
Some(model),
83+
Some(&format!("Panel {}", idx + 1)),
84+
)
85+
.ok()
86+
.flatten()
87+
})
88+
};
89+
if let Some(p) = result {
6890
panel_paths.lock().unwrap()[idx] = Some(p);
6991
}
7092
});
@@ -91,21 +113,26 @@ pub fn run(
91113
gen_model: Option<&str>,
92114
batch: bool,
93115
) -> Result<Option<PathBuf>> {
94-
let gemini_key = cfg
95-
.gemini_key()
96-
.ok_or_else(|| eyre::eyre!("Gemini API key required"))?;
97-
let gemini = GeminiClient::new(gemini_key);
116+
let gemini = cfg.gemini_key().map(GeminiClient::new);
117+
let openai = cfg.openai_key().map(OpenAIImageClient::new);
118+
119+
let model = gen_model.unwrap_or(cfg.gen_model());
120+
if model.starts_with("gpt-image") && openai.is_none() {
121+
eyre::bail!("OpenAI API key required for gpt-image models");
122+
}
123+
if !model.starts_with("gpt-image") && gemini.is_none() {
124+
eyre::bail!("Gemini API key required");
125+
}
98126

99127
std::fs::create_dir_all(out_dir)?;
100128
let total = panels.len();
101-
let model = gen_model.unwrap_or(cfg.gen_model());
102129
let panel_aspect = if layout == "vertical" { "16:9" } else { "1:1" };
103130

104131
eprintln!("Generating {total}-panel comic ({layout})...");
105132

106133
// Phase 1: Generate panels
107-
let mut panel_paths_vec: Vec<Option<PathBuf>> = if batch {
108-
// Batch API path
134+
let mut panel_paths_vec: Vec<Option<PathBuf>> = if batch && !model.starts_with("gpt-image") {
135+
// Batch API path (Gemini only)
109136
let requests: Vec<BatchImageRequest> = panels
110137
.iter()
111138
.enumerate()
@@ -132,12 +159,13 @@ pub fn run(
132159
}
133160
})
134161
.collect();
135-
match gemini.batch_gen_images(requests) {
162+
match gemini.as_ref().unwrap().batch_gen_images(requests) {
136163
Ok(results) => results,
137164
Err(e) => {
138165
eprintln!("Batch failed ({e}), falling back to parallel sync...");
139166
gen_panels_sync(
140167
&gemini,
168+
&openai,
141169
out_dir,
142170
panels,
143171
style,
@@ -152,6 +180,7 @@ pub fn run(
152180
} else {
153181
gen_panels_sync(
154182
&gemini,
183+
&openai,
155184
out_dir,
156185
panels,
157186
style,

0 commit comments

Comments
 (0)