import gradio as gr import torch from diffusers import StableDiffusionPipeline from .utils.schedulers import SCHEDULER_LIST, get_scheduler_list from .utils.prompt2prompt import generate from .utils.device import get_device from .download import get_share_js, community_icon_html, loading_icon_html, CSS #--- create a download button that takes the output image from gradio and downloads it TEXT2IMG_MODEL_LIST = { "OpenJourney v4" : "prompthero/openjourney-v4", "StableDiffusion 1.5" : "runwayml/stable-diffusion-v1-5", "StableDiffusion 2.1" : "stabilityai/stable-diffusion-2-1", "DreamLike 1.0" : "dreamlike-art/dreamlike-diffusion-1.0", "DreamLike 2.0" : "dreamlike-art/dreamlike-photoreal-2.0", "DreamShaper" : "Lykon/DreamShaper", "NeverEnding-Dream" : "Lykon/NeverEnding-Dream" } class StableDiffusionText2ImageGenerator: def __init__(self): self.pipe = None def load_model( self, model_path, scheduler ): model_path = TEXT2IMG_MODEL_LIST[model_path] if self.pipe is None: self.pipe = StableDiffusionPipeline.from_pretrained( model_path, safety_checker=None, torch_dtype=torch.float32 ) device = get_device() self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler) self.pipe.to(device) self.pipe.enable_attention_slicing() return self.pipe def generate_image( self, model_path: str, prompt: str, negative_prompt: str, scheduler: str, guidance_scale: int, num_inference_step: int, height: int, width: int, seed_generator=0, ): print("model_path", model_path) print("prompt", prompt) print("negative_prompt", negative_prompt) print("num_images_per_prompt", 1) print("scheduler", scheduler) print("guidance_scale", guidance_scale) print("num_inference_step", num_inference_step) print("height", height) print("width", width) print("seed_generator", seed_generator) pipe = self.load_model( model_path=model_path, scheduler=scheduler, ) if seed_generator == 0: random_seed = torch.randint(0, 1000000, (1,)) generator = torch.manual_seed(random_seed) else: generator = torch.manual_seed(seed_generator) images = pipe( prompt=prompt, height=height, width=width, negative_prompt=negative_prompt, num_images_per_prompt=1, num_inference_steps=num_inference_step, guidance_scale=guidance_scale, generator=generator, ).images return images def app(username : str = "admin"): demo = gr.Blocks(css = CSS) with demo: with gr.Row(): with gr.Column(): text2image_prompt = gr.Textbox( lines=1, show_label=False, elem_id="prompt-text-input", value='', placeholder="Prompt, keywords that describe your image" ) text2image_negative_prompt = gr.Textbox( lines=1, show_label=False, elem_id = "negative-prompt-text-input", value='', placeholder="Negative Prompt, keywords that describe what you don't want in your image", ) # add button for generating a prompt from the prompt text2image_prompt_generate_button = gr.Button( label="Generate Prompt", type="primary", align="center", value = "Generate Prompt" ) # show a text box with the generated prompt text2image_prompt_generated_prompt = gr.Textbox( lines=1, placeholder="Generated Prompt", show_label=False, info="Auto generated prompts for inspiration.", ) text2image_model_path = gr.Dropdown( choices=list(TEXT2IMG_MODEL_LIST.keys()), value=list(TEXT2IMG_MODEL_LIST.keys())[0], label="Text2Image Model Selection", elem_id="model-dropdown", info="Select the model you want to use for text2image generation." ) text2image_scheduler = gr.Dropdown( choices=SCHEDULER_LIST, value=SCHEDULER_LIST[0], label="Scheduler", elem_id="scheduler-dropdown", info="Scheduler list for models. Different schdulers result in different outputs." ) text2image_size = gr.Slider( minimum=128, maximum=1280, step=32, value=768, label="Image Size", elem_id="image-size-slider", info = "Image size determines the height and width of the generated image. Higher the value, better the quality however slower the computation." ) text2image_seed_generator = gr.Slider( label="Seed(0 for random)", minimum=0, maximum=1000000, value=0, elem_id="seed-slider", info="Set the seed to a specific value to reproduce the results." ) text2image_guidance_scale = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale", elem_id = "guidance-scale-slider", info = "Guidance scale determines how much the prompt will affect the image. Higher the value, more the effect." ) text2image_num_inference_step = gr.Slider( minimum=1, maximum=100, step=1, value=50, label="Num Inference Step", elem_id = "num-inference-step-slider", info = "Number of inference step determines the quality of the image. Higher the number, better the quality." ) text2image_predict = gr.Button(value="Generate image") with gr.Column(): output_image = gr.Gallery( label="Generated images", show_label=False, elem_id="gallery", ).style(grid=(1, 2), height='auto') with gr.Group(elem_id="container-advanced-btns"): with gr.Group(elem_id="share-btn-container"): community_icon = gr.HTML(community_icon_html) loading_icon = gr.HTML(loading_icon_html) share_button = gr.Button("Save artwork", elem_id="share-btn") gr.HTML( """

Text2Image Models

Text to image models will generate an image guided by the prompt that is provided

A prompt should be specified with keywords that describe the image you want to generate.

Negative prompt can be used to specify keywords that you don't want in your image such as "blood" or "violence".

Example prompt: "A painting of a cat sitting on a chair, fantasy themed, starry background"

Stable Diffusion 1.5 & 2.1: Default model for many tasks.

OpenJourney v4: Generates fantasy themed images similar to the Midjourney model.

Dreamlike Photoreal 1.0 & 2.0 is SD 1.5 that generates realistic images.

""" ) text2image_predict.click( fn=StableDiffusionText2ImageGenerator().generate_image, inputs=[ text2image_model_path, text2image_prompt, text2image_negative_prompt, text2image_scheduler, text2image_guidance_scale, text2image_num_inference_step, text2image_size, text2image_size, text2image_seed_generator, ], outputs=output_image, ) text2image_prompt_generate_button.click( fn=generate, inputs=[text2image_prompt], outputs=[text2image_prompt_generated_prompt], ) # share_button.click( # None, # [], # [], # _js=get_share_js(), # ) # autoclik the share button return demo