rewrite app as blocks and use image generation

app.py

@@ -1,4 +1,3 @@
-import os
 import gradio as gr
 
 import torch
@@ -7,60 +6,73 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 
-summarizer = pipeline("summarization")
-model_id = "runwayml/stable-diffusion-v1-5"
-
-SAVED_CHECKPOINT = 'mikegarts/distilgpt2-lotr'
-MIN_WORDS = 120
-
 READ_TOKEN = os.environ.get('HF_ACCESS_TOKEN', None)
 
-def get_image_pipe():
+model_id = "runwayml/stable-diffusion-v1-5"
+# model_id = "CompVis/stable-diffusion-v1-4"
+
+has_cuda = torch.cuda.is_available()
+device = "cpu"
+if has_cuda:
     pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=READ_TOKEN)
-    pipe.to('cuda')
-    return pipe
+    device = "cuda"
+else:
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", use_auth_token=READ_TOKEN)
+    device = "cpu"
+    
+pipe.to(device)
+def safety_checker(images, clip_input):
+    return images, False
+pipe.safety_checker = safety_checker
 
-def get_model():
-    model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
-    tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
-    return model, tokenizer
+SAVED_CHECKPOINT = 'mikegarts/distilgpt2-lotr'
+model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
+tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
 
+summarizer = pipeline("summarization")
 
-def generate(prompt):
-    model, tokenizer = get_model()
+def break_until_dot(txt):
+    return txt.rsplit('.', 1)[0] + '.'
 
+def generate(prompt):
     input_context = prompt
     input_ids = tokenizer.encode(input_context, return_tensors="pt").to(model.device)
 
     outputs = model.generate(
         input_ids=input_ids, 
-        max_length=100, 
+        max_length=180, 
         temperature=0.7, 
         num_return_sequences=3, 
         do_sample=True
     )
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return break_until_dot(decoded)
 
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).rsplit('.', 1)[0] + '.'
-
-def make_image(prompt):
-    pipe = get_image_pipe()
-    image = pipe(prompt)["sample"][0] 
+def generate_image(prompt, inference_steps):
+    prompt = prompt + ', masterpiece charcoal pencil art lord of the rings illustration'
+    img = pipe(prompt, height=512, width=512, num_inference_steps=inference_steps)
+    return img.images[0]
 
-def predict(prompt):
+def generate_story(prompt):
     story = generate(prompt=prompt)
-    summary = summarizer(story, min_length=5, max_length=20)[0]['summary_text']
-    image = make_image(summary)
-    return story, summarizer(story, min_length=5, max_length=20), image
-
-
-title = "Lord of the rings app"
-description = """A Lord of the rings insired app that combines text and image generation"""
-
-gr.Interface(
-    fn=predict,
-    inputs="textbox",
-    outputs=["text", "text", "image"],
-    title=title,
-    description=description,
-    examples=[["My new adventure would be"], ["Then I a hobbit appeared"], ["Frodo told me"]]
-).launch(debug=True)
+    summary = summarizer(story, min_length=5, max_length=15)[0]['summary_text']
+    summary = break_until_dot(summary)
+    return story, summary, gr.update(visible=True)
+
+with gr.Blocks() as demo:
+    title = gr.Markdown('## Lord of the rings app')
+    description = gr.Markdown('### A Lord of the rings insired app that combines text and image generation')
+    prompt = gr.Textbox(label="Your prompt", value="And then the hobbit said")
+    story = gr.Textbox(label="Your story")
+    summary = gr.Textbox(label="Summary")
+    
+    bt_make_text = gr.Button("Generate text")
+    bt_make_image = gr.Button("Generate and image (takes about 10-15 minutes on CPU)", visible=False)
+    
+    image = gr.Image(label='Illustration for your story')
+    inference_steps = gr.Slider(5, 35, value=15, step=1, label="Num inference steps (more steps makes a better image but takes more time)")
+    
+    bt_make_text.click(fn=generate_story, inputs=prompt, outputs=[story, summary, bt_make_image])
+    bt_make_image.click(fn=generate_image, inputs=[summary, inference_steps], outputs=image)
+
+demo.launch(share=True)