| 43 | } |
| 44 | |
| 45 | pub async fn generate_image<M: Model>( |
| 46 | state: web::Data<Arc<RwLock<Master<M>>>>, |
| 47 | req: HttpRequest, |
| 48 | image_request: web::Json<ImageRequest>, |
| 49 | ) -> impl Responder { |
| 50 | let client = req |
| 51 | .peer_addr() |
| 52 | .map(|a| a.to_string()) |
| 53 | .unwrap_or_else(|| "unknown".to_string()); |
| 54 | |
| 55 | log::info!("starting generating image for {} ...", &client); |
| 56 | |
| 57 | let mut master = state.write().await; |
| 58 | |
| 59 | if !master.model.as_ref().is_some_and(|m| m.output_modality() == OutputModality::Image) { |
| 60 | return HttpResponse::NotFound() |
| 61 | .json(serde_json::json!({"error": "No image model loaded"})); |
| 62 | } |
| 63 | |
| 64 | let result_pngs: Arc<Mutex<Vec<Vec<u8>>>> = Arc::new(Mutex::new(Vec::new())); |
| 65 | let result_pngs_cloned = Arc::clone(&result_pngs); |
| 66 | |
| 67 | if let Err(e) = master |
| 68 | .generate_image(image_request.image_args.clone(), move |images| { |
| 69 | let pngs: Vec<Vec<u8>> = images.iter().map(encode_png).collect(); |
| 70 | let mut locked = result_pngs_cloned.lock().expect("Error acquiring lock"); |
| 71 | locked.extend(pngs); |
| 72 | }) |
| 73 | .await |
| 74 | { |
| 75 | return HttpResponse::InternalServerError() |
| 76 | .json(serde_json::json!({"error": format!("{e}")})); |
| 77 | } |
| 78 | |
| 79 | let locked = result_pngs.lock().expect("Error acquiring lock"); |
| 80 | |
| 81 | if image_request.response_format == "png" { |
| 82 | // Return raw PNG bytes for the first image |
| 83 | if let Some(png) = locked.first() { |
| 84 | HttpResponse::Ok() |
| 85 | .content_type("image/png") |
| 86 | .body(png.clone()) |
| 87 | } else { |
| 88 | HttpResponse::InternalServerError() |
| 89 | .json(serde_json::json!({"error": "No image generated"})) |
| 90 | } |
| 91 | } else { |
| 92 | // Default: b64_json (backwards-compatible) |
| 93 | let images: Vec<String> = locked |
| 94 | .iter() |
| 95 | .map(|png| general_purpose::STANDARD.encode(png)) |
| 96 | .collect(); |
| 97 | HttpResponse::Ok().json(ImageResponse { images }) |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // ── OpenAI-compatible image endpoint (/v1/images/generations) ── |
| 102 | |