shivi commited on
Commit
4f8e447
1 Parent(s): 40e1f85

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import from_pretrained_keras
2
+ import keras_cv
3
+
4
+
5
+ def generate_images(prompt: str, num_imgs_to_gen: int):
6
+ """
7
+ This function is used to generate images using our fine-tuned keras dreambooth stable diffusion model.
8
+ Args:
9
+ prompt (str): The text input given by the user based on which images will be generated.
10
+ num_imgs_to_gen (int): The number of images to be generated using given prompt.
11
+ Returns:
12
+ all_images (List): List of images that were generated using the model
13
+ """
14
+ dreambooth_model = keras_cv.models.StableDiffusion(
15
+ img_width=resolution, img_height=resolution, jit_compile=True,
16
+ )
17
+ loaded_diffusion_model = from_pretrained_keras("shivi/dreambooth_diffusion_model")
18
+ dreambooth_model._diffusion_model = loaded_diffusion_model
19
+
20
+ generated_img = dreambooth_model.text_to_image(
21
+ prompt, batch_size=num_imgs_to_gen
22
+ )
23
+
24
+ return generated_img
25
+
26
+
27
+ with gr.Blocks() as demo:
28
+ gr.HTML("<h2 style=\"font-size: 2em; font-weight: bold\" align=\"center\">Keras Dreambooth - Pink Floyd Division Bell Demo</h2>")
29
+
30
+ gr.Markdown("This model has been fine tuned to learn the concept of Division Bell from Pink Floyd's famous album `The Division Bell`")
31
+ with gr.Row():
32
+ with gr.Column():
33
+ prompt = gr.Textbox(label="prompt")
34
+ samples = gr.Slider(label="No. of Images",value=1)
35
+ run = gr.Button(value="Run")
36
+ with gr.Column():
37
+ gallery = gr.Gallery(show_label=False)
38
+
39
+ run.click(generate_images, inputs=[prompt,samples], outputs=gallery)
40
+
41
+ gr.Examples([["A flower vase inspired by pink floyd division bell album cover", 1],
42
+ ["A pendant in the style of pink floyd division bell album cover",1],
43
+ ["A human skull inspired by pink floyd division bell album cover ",1],
44
+ ["picture of pink floyd division bell album cover with a starry night on Mars theme", 1]],
45
+ [prompt,samples,model_choice], gallery, generate_images, cache_examples=True)
46
+
47
+ gr.Markdown("Some fun examples I created while playing with the model:")
48
+ gr.Makdown("![Example 1](sample_outputs/1.png)")
49
+
50
+ demo.launch(debug=True)