Spaces:
Runtime error
Runtime error
| 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) |