Spaces:
Running
Running
| import os | |
| import gradio as gr | |
| import base64 | |
| from random import randint | |
| from all_models import models | |
| from io import BytesIO | |
| from PIL import Image | |
| css_code = os.getenv("DazDinGo_CSS") | |
| def load_fn(models): | |
| global models_load | |
| models_load = {} | |
| for model in models: | |
| if model not in models_load.keys(): | |
| try: | |
| m = gr.load(f'models/{model}') | |
| except Exception as error: | |
| m = gr.Interface(lambda txt: None, ['text'], ['image']) | |
| models_load.update({model: m}) | |
| load_fn(models) | |
| def gen_fn(model_str, prompt): | |
| if model_str == 'NA': | |
| return None, "" | |
| noise = str(randint(0, 4294967296)) | |
| klir = '| ultra detail, ultra elaboration, ultra quality, perfect' | |
| image = models_load[model_str](f'{prompt} {klir} {noise}') | |
| # Convert to base64 | |
| buffered = BytesIO() | |
| if isinstance(image, str): # if it's a file path | |
| img = Image.open(image) | |
| img.save(buffered, format="JPEG") | |
| else: # if it's a PIL Image | |
| image.save(buffered, format="JPEG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode() | |
| base64_code = f"data:image/jpeg;base64,{img_str}" | |
| return image, base64_code | |
| def make_me(): | |
| with gr.Row(): | |
| with gr.Column(scale=4): | |
| txt_input = gr.Textbox(label='Your prompt:', lines=4, container=False, | |
| elem_id="custom_textbox", placeholder="Prompt", height=250) | |
| with gr.Column(scale=1): | |
| gen_button = gr.Button('Generate image', elem_id="custom_gen_button") | |
| stop_button = gr.Button('Stop', variant='secondary', interactive=False, | |
| elem_id="custom_stop_button") | |
| def on_generate_click(): | |
| return gr.Button('Generate image', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button") | |
| def on_stop_click(): | |
| return gr.Button('Generate image', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button") | |
| gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button]) | |
| stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button]) | |
| with gr.Row(): | |
| with gr.Column(): | |
| model_dropdown = gr.Dropdown(models, label="Select Model", | |
| value=models[0] if models else None) | |
| output_image = gr.Image(label="Generated Image", width=512, | |
| max_height=768, elem_id="custom_image", | |
| show_label=True, interactive=False, | |
| show_share_button=False) | |
| # Base64 output | |
| base64_output = gr.Textbox(label="Base64 Code", interactive=True, | |
| placeholder="Base64 code will appear here after generation", | |
| lines=4) | |
| gen_event = gen_button.click(gen_fn, [model_dropdown, txt_input], | |
| [output_image, base64_output]) | |
| stop_button.click(on_stop_click, inputs=None, | |
| outputs=[gen_button, stop_button], cancels=[gen_event]) | |
| with gr.Row(): | |
| gr.HTML("") | |
| with gr.Blocks(css=css_code) as demo: | |
| make_me() | |
| demo.queue(concurrency_count=50) | |
| demo.launch() |