feat: add getCurrentStep function

example/App.tsx

@@ -6,8 +6,10 @@ import * as ExpoStableDiffusion from "expo-stable-diffusion";
 
 const MODEL_PATH = FileSystem.documentDirectory + "Model/stable-diffusion-2-1";
 const SAVE_PATH = FileSystem.documentDirectory + "image.jpeg";
+const STEP_COUNT = 25;
 
 export default function App() {
+  const [currentStep, setCurrentStep] = React.useState(0);
   React.useEffect(() => {
     (async () => {
       Alert.alert(`Loading Model: ${MODEL_PATH}`);
@@ -16,19 +18,31 @@ export default function App() {
 
       Alert.alert("Model Loaded, Generating Images!");
 
+      setCurrentStep(0);
       await ExpoStableDiffusion.generateImage({
         prompt: "a cat coding at night",
-        stepCount: 25,
+        stepCount: STEP_COUNT,
         savePath: SAVE_PATH,
       });
 
       Alert.alert(`Image Generated: ${SAVE_PATH}`);
     })();
+
+    const timer = setInterval(() => {
+      setCurrentStep(ExpoStableDiffusion.getCurrentStep());
+    }, 1000);
+
+    return () => {
+      clearInterval(timer);
+    };
   }, []);
 
   return (
     <View style={styles.container}>
       <Text>Testing Expo Stable Diffusion</Text>
+      <Text>
+        Progress: {currentStep}/{STEP_COUNT}
+      </Text>
     </View>
   );
 }

ios/ExpoStableDiffusionModule.swift

@@ -3,13 +3,16 @@ import CoreML
 
 public class ExpoStableDiffusionModule: Module {
     private var pipeline: StableDiffusionPipeline? = nil
+    private var currentStep: Int = 0
 
     public func definition() -> ModuleDefinition {
         Name("ExpoStableDiffusion")
 
         AsyncFunction("loadModel", loadModel)
 
         AsyncFunction("generateImage", generateImage)
+
+        Function("getCurrentStep", getCurrentStep)
     }
 
     private func loadModel(modelPath: URL) throws {
@@ -29,11 +32,13 @@ public class ExpoStableDiffusionModule: Module {
         var config = StableDiffusionPipeline.Configuration(prompt: prompt)
         config.schedulerType = .dpmSolverMultistepScheduler
         config.stepCount = stepCount!
+        self.currentStep = 0
 
         print("Generating Images with the following Config:", config)
 
         let image = try self.pipeline!.generateImages(configuration: config, progressHandler: { progress in
             print("Current Step: \(progress.step)")
+            self.currentStep = progress.step
             return true
         }).first
 
@@ -45,4 +50,8 @@ public class ExpoStableDiffusionModule: Module {
 
         print("Image Generated at: \(savePath)")
     }
+
+    private func getCurrentStep() -> Int {
+        return self.currentStep
+    }
 }

src/index.ts

@@ -16,3 +16,7 @@ export async function generateImage({
     savePath
   );
 }
+
+export function getCurrentStep(): number {
+  return ExpoStableDiffusionModule.getCurrentStep();
+}