Spaces:
Runtime error
Runtime error
File size: 3,107 Bytes
7547682 71803ee 7547682 bf00a84 7547682 bf00a84 7547682 6e36368 7547682 |
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 |
import os
import sys
import sys
import torch
import gradio as gr
from stable_diffusion import StableDiffusion
from utils import invert_loss, get_style_embeddings, show_images
torch.manual_seed(1)
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
def gradio_interface(prompt:str, style_name:str):
stable_diffuser = StableDiffusion(torch_device)
#outputs_1 = []
#outputs_2 = []
seed_values = [1,2,3,4,5]
custom_loss_scale = 150.0
num_styles = len(style_files)
style_file = style_files[style_name]
style_token_embedding = get_style_embeddings(style_file) if style_file is not None else None
this_generated_img_1 = stable_diffuser.generate_image_with_custom_style(prompt,
style_token_embedding = style_token_embedding,
random_seed = 42,
custom_loss_fn = None)
#outputs_1.append(this_generated_img_1)
this_generated_img_2 = stable_diffuser.generate_image_with_custom_style(prompt,
style_token_embedding = style_token_embedding,
random_seed = 42,
custom_loss_fn = invert_loss,
custom_loss_scale = custom_loss_scale)
#outputs_2.append(this_generated_img_2)
return this_generated_img_1, this_generated_img_2
style_files = {'No_style': None,
'watercolor': 'learned_embeds_watercolor.bin',
'strip': 'learned_embeds_strip_style.bin',
'oil_paint': 'learned_embeds_oil_paint.bin',
'kaleido': 'learned_embeds_kaleido.bin',
'doodle': 'learned_embeds_doodle.bin'}
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
prompt_examples = [["A person wearing a hat and smoking a cigar", 'watercolor']]
#["An oil painting of a bear playing guitar"]]
#["An oil painting of a lion eating pizza and drinking coke"]]
# Define Interface
description = 'A Stable Diffusion based Generative AI tool to generate images. Also generates grayscale images.'
title = 'Image Generation using Stable Diffusion with styles.'
style_names = ['No_style','watercolor','strip','oil_paint','kaleido','doodle']
demo = gr.Interface(gradio_interface,
inputs = [gr.Textbox('A deer crossing the street', label="Text Prompt"),
gr.Dropdown(style_names,value='watercolor',label="Display Style")],
outputs = [gr.Image(shape=(512, 512),label='Generated Image').style(width=512, height=512),
gr.Image(shape=(512, 512),label='Generated Images with custom loss').style(width=512, height=512)
],
examples=prompt_examples,
title = title,
description = description
)
demo.launch(debug=True) |