feat: Add text to image pipeline

.cspell-wordlist.txt

@@ -66,3 +66,11 @@ softmax
 logit
 logits
 probs
+unet
+Unet
+VPRED
+timesteps
+Timesteps
+denoises
+denoise
+denoising

apps/computer-vision/app/_layout.tsx

@@ -100,6 +100,14 @@ export default function _layout() {
             headerTitleStyle: { color: ColorPalette.primary },
           }}
         />
+        <Drawer.Screen
+          name="text_to_image/index"
+          options={{
+            drawerLabel: 'Image Generation',
+            title: 'Image Generation',
+            headerTitleStyle: { color: ColorPalette.primary },
+          }}
+        />
         <Drawer.Screen
           name="index"
           options={{

apps/computer-vision/app/index.tsx

@@ -47,6 +47,12 @@ export default function Home() {
         >
           <Text style={styles.buttonText}>Style Transfer</Text>
         </TouchableOpacity>
+        <TouchableOpacity
+          style={styles.button}
+          onPress={() => router.navigate('text_to_image/')}
+        >
+          <Text style={styles.buttonText}>Image Generation</Text>
+        </TouchableOpacity>
       </View>
     </View>
   );

apps/computer-vision/app/text_to_image/index.tsx

@@ -0,0 +1,183 @@
+import {
+  View,
+  StyleSheet,
+  Text,
+  Image,
+  Keyboard,
+  TouchableWithoutFeedback,
+} from 'react-native';
+import React, { useContext, useEffect, useState } from 'react';
+import Spinner from 'react-native-loading-spinner-overlay';
+import { useTextToImage, BK_SDM_TINY_VPRED_256 } from 'react-native-executorch';
+import { GeneratingContext } from '../../context';
+import ColorPalette from '../../colors';
+import ProgressBar from '../../components/ProgressBar';
+import { BottomBarWithTextInput } from '../../components/BottomBarWithTextInput';
+
+export default function TextToImageScreen() {
+  const [inferenceStepIdx, setInferenceStepIdx] = useState<number>(0);
+  const [imageTitle, setImageTitle] = useState<string | null>(null);
+  const [image, setImage] = useState<string | null>(null);
+  const [steps, setSteps] = useState<number>(10);
+  const [showTextInput, setShowTextInput] = useState(false);
+  const [keyboardVisible, setKeyboardVisible] = useState(false);
+
+  const imageSize = 224;
+  const model = useTextToImage({
+    model: BK_SDM_TINY_VPRED_256,
+    inferenceCallback: (x) => setInferenceStepIdx(x),
+  });
+
+  const { setGlobalGenerating } = useContext(GeneratingContext);
+
+  useEffect(() => {
+    setGlobalGenerating(model.isGenerating);
+  }, [model.isGenerating, setGlobalGenerating]);
+
+  useEffect(() => {
+    const showSub = Keyboard.addListener('keyboardDidShow', () => {
+      setKeyboardVisible(true);
+    });
+    const hideSub = Keyboard.addListener('keyboardDidHide', () => {
+      setKeyboardVisible(false);
+    });
+    return () => {
+      showSub.remove();
+      hideSub.remove();
+    };
+  }, []);
+
+  const runForward = async (input: string, numSteps: number) => {
+    if (!input || !input.trim()) return;
+    const prevImageTitle = imageTitle;
+    setImageTitle(input);
+    setSteps(numSteps);
+    try {
+      const output = await model.generate(input, imageSize, steps);
+      if (!output.length) {
+        setImageTitle(prevImageTitle);
+        return;
+      }
+      setImage(output);
+    } catch (e) {
+      console.error(e);
+      setImageTitle(null);
+    } finally {
+      setInferenceStepIdx(0);
+    }
+  };
+
+  if (!model.isReady) {
+    // TODO: Update when #614 merged
+    return (
+      <Spinner
+        visible={!model.isReady}
+        textContent={`Loading the model ${(model.downloadProgress * 100).toFixed(0)} %`}
+      />
+    );
+  }
+
+  return (
+    <TouchableWithoutFeedback
+      onPress={() => {
+        Keyboard.dismiss();
+        setShowTextInput(false);
+      }}
+    >
+      <View style={styles.container}>
+        {keyboardVisible && <View style={styles.overlay} />}
+
+        <View style={styles.titleContainer}>
+          {imageTitle && <Text style={styles.titleText}>{imageTitle}</Text>}
+        </View>
+
+        {model.isGenerating ? (
+          <View style={styles.progressContainer}>
+            <Text style={styles.text}>Generating...</Text>
+            <ProgressBar numSteps={steps} currentStep={inferenceStepIdx} />
+          </View>
+        ) : (
+          <View style={styles.imageContainer}>
+            {image?.length ? (
+              <Image
+                style={styles.image}
+                source={{ uri: `data:image/png;base64,${image}` }}
+              />
+            ) : (
+              <Image
+                style={styles.image}
+                source={require('../../assets/icons/executorch_logo.png')}
+              />
+            )}
+          </View>
+        )}
+
+        <View style={styles.bottomContainer}>
+          <BottomBarWithTextInput
+            runModel={runForward}
+            stopModel={model.interrupt}
+            isGenerating={model.isGenerating}
+            isReady={model.isReady}
+            showTextInput={showTextInput}
+            setShowTextInput={setShowTextInput}
+            keyboardVisible={keyboardVisible}
+          />
+        </View>
+      </View>
+    </TouchableWithoutFeedback>
+  );
+}
+
+const styles = StyleSheet.create({
+  container: {
+    flex: 1,
+    width: '100%',
+    alignItems: 'center',
+  },
+  overlay: {
+    ...StyleSheet.absoluteFillObject,
+    backgroundColor: 'rgba(0,0,0,0.65)',
+    zIndex: 5,
+  },
+  titleContainer: {
+    alignItems: 'center',
+    marginTop: 20,
+  },
+  titleText: {
+    color: ColorPalette.primary,
+    fontSize: 20,
+    fontWeight: 'bold',
+    marginBottom: 12,
+    textAlign: 'center',
+  },
+  text: {
+    fontSize: 16,
+    color: '#000',
+  },
+  imageContainer: {
+    flex: 1,
+    position: 'absolute',
+    top: 100,
+    alignItems: 'center',
+    justifyContent: 'center',
+  },
+  image: {
+    width: 256,
+    height: 256,
+    marginVertical: 30,
+    resizeMode: 'contain',
+  },
+  progressContainer: {
+    flex: 1,
+    justifyContent: 'center',
+    alignItems: 'center',
+  },
+  bottomContainer: {
+    flex: 1,
+    width: '90%',
+    position: 'absolute',
+    bottom: 0,
+    marginBottom: 25,
+    zIndex: 10,
+  },
+});

apps/computer-vision/components/BottomBarWithTextInput.tsx

@@ -0,0 +1,165 @@
+import React, { useState } from 'react';
+import {
+  View,
+  Text,
+  TextInput,
+  TouchableOpacity,
+  StyleSheet,
+  KeyboardAvoidingView,
+  Platform,
+} from 'react-native';
+import { Ionicons } from '@expo/vector-icons';
+import ColorPalette from '../colors';
+
+interface BottomBarProps {
+  runModel: (input: string, numSteps: number) => void;
+  stopModel: () => void;
+  isGenerating?: boolean;
+  isReady?: boolean;
+  showTextInput: boolean;
+  setShowTextInput: React.Dispatch<React.SetStateAction<boolean>>;
+  keyboardVisible: boolean;
+}
+
+export const BottomBarWithTextInput = ({
+  runModel,
+  stopModel,
+  isGenerating,
+  isReady,
+  showTextInput,
+  setShowTextInput,
+  keyboardVisible,
+}: BottomBarProps) => {
+  const [input, setInput] = useState('');
+  const [numSteps, setNumSteps] = useState(10);
+
+  const decreaseSteps = () => setNumSteps((prev) => Math.max(5, prev - 5));
+  const increaseSteps = () => setNumSteps((prev) => Math.min(50, prev + 5));
+
+  if (!showTextInput) {
+    if (isGenerating) {
+      return (
+        <TouchableOpacity
+          style={styles.button}
+          onPress={stopModel}
+          disabled={!isReady}
+        >
+          <Text style={styles.buttonText}>Stop model</Text>
+        </TouchableOpacity>
+      );
+    } else {
+      return (
+        <TouchableOpacity
+          style={styles.button}
+          onPress={() => setShowTextInput(true)}
+          disabled={!isReady}
+        >
+          <Text style={styles.buttonText}>Run model</Text>
+        </TouchableOpacity>
+      );
+    }
+  }
+
+  return (
+    <KeyboardAvoidingView
+      style={styles.container}
+      collapsable={false}
+      behavior={Platform.OS === 'ios' ? 'padding' : undefined}
+      keyboardVerticalOffset={Platform.OS === 'ios' ? 120 : 40}
+    >
+      <View style={styles.inputContainer}>
+        <TextInput
+          style={styles.input}
+          placeholder="Enter prompt..."
+          value={input}
+          onChangeText={setInput}
+        />
+        <TouchableOpacity
+          style={[styles.button, styles.iconButton]}
+          onPress={() => {
+            setShowTextInput(false);
+            setInput('');
+            runModel(input, numSteps);
+          }}
+          disabled={!isReady || isGenerating}
+        >
+          <Ionicons name="send" size={20} color="#fff" />
+        </TouchableOpacity>
+      </View>
+
+      <View style={styles.stepsContainer}>
+        <Text style={[styles.text, keyboardVisible && styles.textWhite]}>
+          Steps: {numSteps}
+        </Text>
+        <View style={styles.stepsButtons}>
+          <TouchableOpacity
+            style={[styles.button, styles.iconButton]}
+            onPress={decreaseSteps}
+          >
+            <Text style={styles.buttonText}>-</Text>
+          </TouchableOpacity>
+          <TouchableOpacity
+            style={[styles.button, styles.iconButton]}
+            onPress={increaseSteps}
+          >
+            <Text style={styles.buttonText}>+</Text>
+          </TouchableOpacity>
+        </View>
+      </View>
+    </KeyboardAvoidingView>
+  );
+};
+
+const styles = StyleSheet.create({
+  container: {
+    alignItems: 'center',
+  },
+  inputContainer: {
+    flexDirection: 'row',
+    alignItems: 'center',
+    justifyContent: 'center',
+  },
+  input: {
+    flex: 1,
+    borderRadius: 6,
+    padding: 8,
+    marginRight: 8,
+    backgroundColor: '#fff',
+    color: '#000',
+  },
+  stepsContainer: {
+    width: '100%',
+    flexDirection: 'row',
+    alignItems: 'center',
+    justifyContent: 'space-between',
+    marginTop: 10,
+  },
+  stepsButtons: {
+    flexDirection: 'row',
+  },
+  button: {
+    width: '100%',
+    height: 40,
+    justifyContent: 'center',
+    alignItems: 'center',
+    backgroundColor: ColorPalette.primary,
+    borderRadius: 8,
+  },
+  buttonText: {
+    color: '#fff',
+    fontSize: 16,
+    textAlign: 'center',
+  },
+  iconButton: {
+    marginHorizontal: 5,
+    width: 40,
+  },
+  text: {
+    flex: 1,
+    fontSize: 16,
+    color: '#000',
+  },
+  textWhite: {
+    color: '#fff',
+  },
+});