File size: 2,959 Bytes
34d353a
05d7e61
 
 
 
 
 
 
 
79c9a48
 
3e4033b
 
 
 
34d353a
3e4033b
79c9a48
3e4033b
 
 
 
 
 
 
 
 
05d7e61
3e4033b
 
 
05d7e61
3e4033b
05d7e61
3e4033b
 
05d7e61
3e4033b
05d7e61
 
 
 
 
3e4033b
05d7e61
 
 
 
3e4033b
 
05d7e61
3e4033b
 
 
 
05d7e61
3e4033b
05d7e61
3e4033b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34d353a
3e4033b
 
 
 
34d353a
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import os
import gradio as gr

import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import pipeline
from diffusers import StableDiffusionPipeline

READ_TOKEN = os.environ.get('HF_ACCESS_TOKEN', None)

model_id = "runwayml/stable-diffusion-v1-5"
# model_id = "CompVis/stable-diffusion-v1-4"

has_cuda = torch.cuda.is_available()

if has_cuda:
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=READ_TOKEN)
    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

SAVED_CHECKPOINT = 'mikegarts/distilgpt2-lotr'
model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)

summarizer = pipeline("summarization")

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=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)

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 generate_story(prompt):
    story = generate(prompt=prompt)
    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, 30, 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)

if READ_TOKEN:
    demo.launch()
else:
    demo.launch(share=True, debug=True)