Update index.mjs

index.mjs

@@ -4,6 +4,8 @@ import { createServer } from "http";
 import { Server } from "socket.io";
 import { client } from "@gradio/client";
 import { createRequire } from 'node:module'
+import fs from 'fs';
+import path from 'path';
 const require = createRequire(import.meta.url);
 global.EventSource = require('eventsource');
 
@@ -13,7 +15,6 @@ app.use(express.static('public'));
 const io = new Server(httpServer, { /* options */ });
 
 io.on("connection", (socket) => {
-    
     console.log("new socket connection");
     
     socket.on("ask_api", (client_data) => {
@@ -21,49 +22,52 @@ io.on("connection", (socket) => {
         console.log("trying to reach api");   
         asyncAPICall(client_data, socket)
     });
-
 });
 
-async function test_servers(){
-  try{
-    const grapi_test = await client("https://gradio-hello-world.hf.space");
-    const apitest_result = await grapi_test.predict("/predict", [
-      "John",
-    ]);
-    console.log(apitest_result);
-  }
-  catch(e){
-    console.log(e)
-  }
-  
-}
-
-//test_servers()
-
 async function asyncAPICall(data, socket) {
+    const grapi = await client("fffiloni/mndrm-call");
+    try {
+        const api_result = await grapi.predict("/infer", [
+            data[0], // blob in 'image' Image component		
+            data[1], // string in 'Question' Textbox component
+        ]); 
+        console.log(api_result)
 
-  const grapi = await client("fffiloni/mndrm-call");
-  try{
-      const api_result = await grapi.predict("/infer", [
-          data[0], 	// blob in 'image' Image component		
-          data[1], // string  in 'Question' Textbox component
-      ]); 
-      console.log(api_result)
-      socket.emit("api_response", (api_result.data))
-  }
-  catch(e){
-      console.log(e)
-      socket.emit("api_error", ("ERROR ON API SIDE, SORRY..."))
-  }    
- 
+        // Save the image and provide a download link
+        if(api_result && api_result.data){
+            const filename = createFileName(api_result.data.text);
+            saveImage(data[0], filename);
+            socket.emit("api_response", { ...api_result.data, downloadLink: `/download/${filename}` });
+        }
+    } catch(e) {
+        console.log(e)
+        socket.emit("api_error", ("ERROR ON API SIDE, SORRY..."))
+    }    
 }
 
+// Function to create a filename from the caption and current datetime
+function createFileName(caption) {
+    const datetime = new Date().toISOString().replace(/:/g, '-');
+    return `${caption}-${datetime}.png`;
+}
 
+// Function to save image data to a file
+function saveImage(imageData, filename) {
+    const buffer = Buffer.from(imageData.split(",")[1], 'base64');
+    fs.writeFileSync(path.join(__dirname, 'public', filename), buffer);
+}
+
+// Express route to download the image
+app.get('/download/:filename', (req, res) => {
+    const filename = req.params.filename;
+    const filepath = path.join(__dirname, 'public', filename);
+    res.download(filepath, filename, (err) => {
+        if (err) {
+            console.error(err);
+            res.status(500).send("Error occurred while downloading the file.");
+        }
+    });
+});
 
 httpServer.listen(7860);
 console.log("App running on localhost:7860")
-
-
-
-
-