Update app.py

app.py

@@ -3,6 +3,7 @@ from jax import pmap
 from flax.training.common_utils import shard
 import jax
 import jax.numpy as jnp
+import gradio as gr
 
 from pathlib import Path
 from PIL import Image
@@ -44,12 +45,7 @@ def sd2_inference(pipeline, prompts, params, seed = 42, num_inference_steps = 50
     images = images.reshape((images.shape[0] * images.shape[1], ) + images.shape[-3:])
     images = pipeline.numpy_to_pil(images)
     return images
-def image_grid(imgs, rows, cols, down_sample = 1 ):
-    w,h = imgs[0].size
-    grid = Image.new('RGB', size=(cols*w, rows*h))
-    for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))
-    grid = grid.resize( (grid.size[0]//down_sample, grid.size[1]//down_sample) )
-    return grid
+
 
 
 HF_ACCESS_TOKEN = os.environ["HFAUTH"]
@@ -64,4 +60,60 @@ pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
 )
 
 prompts = ["apple"] * 1
-images = sd2_inference(pipeline, prompts, params, seed = 42, num_inference_steps = 1 )
+
+
+
+
+
+
+
+
+
+def generate_image(dense_class_vector=None, int_index=None, noise_seed_vector=None, truncation=0.4):
+    seed = int(noise_seed_vector.sum().item()) if noise_seed_vector is not None else None
+    noise_vector = truncated_noise_sample(truncation=truncation, batch_size=1, seed=seed)
+    noise_vector = torch.from_numpy(noise_vector)
+    if int_index is not None:
+        class_vector = one_hot_from_int([int_index], batch_size=1)
+        class_vector = torch.from_numpy(class_vector)
+        dense_class_vector = gan_model.embeddings(class_vector)
+    else:
+        if isinstance(dense_class_vector, np.ndarray):
+            dense_class_vector = torch.tensor(dense_class_vector)
+        dense_class_vector = dense_class_vector.view(1, 128)
+
+    input_vector = torch.cat([noise_vector, dense_class_vector], dim=1)
+
+    # Generate an image
+    with torch.no_grad():
+        output = gan_model.generator(input_vector, truncation)
+    output = output.cpu().numpy()
+    output = output.transpose((0, 2, 3, 1))
+    output = ((output + 1.0) / 2.0) * 256
+    output.clip(0, 255, out=output)
+    output = np.asarray(np.uint8(output[0]), dtype=np.uint8)
+    return output
+
+
+
+def text_to_image(text):
+    images = sd2_inference(pipeline, [text], params, seed = 42, num_inference_steps = 5 )
+    img = images[0]
+    return img
+
+examples = ["apple",
+            "banana",
+            "chocolate"]
+
+if __name__ == '__main__':
+    interFace = gr.Interface(fn=text_to_image,
+                             inputs=gr.inputs.Textbox(placeholder="Enter the text to Encode to an image", label="Text "
+                                                                                                           "query",
+                                                      lines=1),
+                             outputs=gr.outputs.Image(type="auto", label="Generated Image"),
+                             verbose=True,
+                             examples=examples,
+                             title="Generate Image from Text",
+                             description="",
+                             theme="huggingface")
+    interFace.launch()