| |
|
|
| from __future__ import annotations |
|
|
| import os |
| import random |
| from typing import Tuple, Optional |
|
|
| import gradio as gr |
| from huggingface_hub import HfApi |
|
|
| from inf import InferencePipeline |
|
|
|
|
| SAMPLE_MODEL_IDS = [ |
| 'thliang01/medieval-knight-sdxl-dora-v1-2', |
| 'lora-library/B-LoRA-teddybear', |
| 'lora-library/B-LoRA-bull', |
| 'lora-library/B-LoRA-wolf_plushie', |
| 'lora-library/B-LoRA-pen_sketch', |
| 'lora-library/B-LoRA-cartoon_line', |
| 'lora-library/B-LoRA-child', |
| 'lora-library/B-LoRA-vase', |
| 'lora-library/B-LoRA-scary_mug', |
| 'lora-library/B-LoRA-statue', |
| 'lora-library/B-LoRA-colorful_teapot', |
| 'lora-library/B-LoRA-grey_sloth_plushie', |
| 'lora-library/B-LoRA-teapot', |
| 'lora-library/B-LoRA-backpack_dog', |
| 'lora-library/B-LoRA-buddha', |
| 'lora-library/B-LoRA-dog6', |
| 'lora-library/B-LoRA-poop_emoji', |
| 'lora-library/B-LoRA-pot', |
| 'lora-library/B-LoRA-fat_bird', |
| 'lora-library/B-LoRA-elephant', |
| 'lora-library/B-LoRA-metal_bird', |
| 'lora-library/B-LoRA-cat', |
| 'lora-library/B-LoRA-dog2', |
| 'lora-library/B-LoRA-drawing1', |
| 'lora-library/B-LoRA-village_oil', |
| 'lora-library/B-LoRA-watercolor', |
| 'lora-library/B-LoRA-house_3d', |
| 'lora-library/B-LoRA-ink_sketch', |
| 'lora-library/B-LoRA-drawing3', |
| 'lora-library/B-LoRA-crayon_drawing', |
| 'lora-library/B-LoRA-kiss', |
| 'lora-library/B-LoRA-drawing4', |
| 'lora-library/B-LoRA-working_cartoon', |
| 'lora-library/B-LoRA-painting', |
| 'lora-library/B-LoRA-drawing2' |
| 'lora-library/B-LoRA-multi-dog2', |
| ] |
| css = """ |
| .gradio-container { |
| max-width: 900px !important; |
| } |
| |
| #title { |
| text-align: center; |
| } |
| |
| #title h1 { |
| font-size: 250%; |
| } |
| |
| .lora-title { |
| background-image: linear-gradient(to right, #314755 0%, #26a0da 51%, #314755 100%); |
| text-align: center; |
| border-radius: 10px; |
| display: block; |
| } |
| |
| .lora-title h2 { |
| color: white !important; |
| } |
| |
| .gr-image { |
| width: 256px; |
| height: 256px; |
| object-fit: contain; |
| margin: auto; |
| } |
| |
| .res-image { |
| object-fit: contain; |
| margin: auto; |
| } |
| |
| .lora-column { |
| border: none; |
| background: none; |
| } |
| .gr-row { |
| align-items: center; |
| justify-content: center; |
| margin-top: 5px; |
| } |
| |
| .svelte-iyf88w { |
| background: none; |
| } |
| """ |
|
|
| def get_choices(hf_token): |
| api = HfApi(token=hf_token) |
| choices = [ |
| info.modelId for info in api.list_models(author='thliang01') |
| ] |
| models_list = ['None'] + SAMPLE_MODEL_IDS + choices |
| return models_list |
|
|
|
|
| def get_image_from_card(card, model_id) -> Optional[str]: |
| try: |
| card_path = f"https://huggingface.co/{model_id}/resolve/main/" |
| widget = card.data.get('widget') |
| if widget is not None or len(widget) > 0: |
| output = widget[0].get('output') |
| if output is not None: |
| url = output.get('url') |
| if url is not None: |
| return card_path + url |
| return None |
| except Exception: |
| return None |
|
|
|
|
| def demo_init(): |
| try: |
| choices = get_choices(app.hf_token) |
| content_blora = random.choice(SAMPLE_MODEL_IDS) |
| style_blora = random.choice(SAMPLE_MODEL_IDS) |
| content_blora_prompt, content_blora_image = app.load_model_info(content_blora) |
| style_blora_prompt, style_blora_image = app.load_model_info(style_blora) |
|
|
| content_lora_model_id = gr.update(choices=choices, value=content_blora) |
| content_prompt = gr.update(value=content_blora_prompt) |
| content_image = gr.update(value=content_blora_image) |
|
|
| style_lora_model_id = gr.update(choices=choices, value=style_blora) |
| style_prompt = gr.update(value=style_blora_prompt) |
| style_image = gr.update(value=style_blora_image) |
|
|
| prompt = gr.update( |
| value=f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style') |
|
|
| return content_lora_model_id, content_prompt, content_image, style_lora_model_id, style_prompt, style_image, prompt |
|
|
| except Exception as e: |
| raise type(e)(f'failed to demo_init, due to: {e}') |
|
|
|
|
| def toggle_column(is_checked): |
| try: |
| return 'None' if is_checked else random.choice(SAMPLE_MODEL_IDS) |
| except Exception as e: |
| raise type(e)(f'failed to toggle_column, due to: {e}') |
|
|
| def handle_prompt_change(content_blora_prompt, style_blora_prompt) -> str: |
| try: |
| if content_blora_prompt and style_blora_prompt: |
| return f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' |
| if content_blora_prompt: |
| return content_blora_prompt |
| if style_blora_prompt: |
| return f'A dog in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' |
|
|
| return '' |
| except Exception as e: |
| raise type(e)(f'failed to handle_prompt_change, due to: {e}') |
|
|
|
|
| class InferenceUtil: |
| def __init__(self, hf_token: str | None): |
| self.hf_token = hf_token |
|
|
| def load_model_info(self, lora_model_id: str) -> Tuple[str, Optional[str]]: |
| try: |
| try: |
| card = InferencePipeline.get_model_card(lora_model_id, |
| self.hf_token) |
| except Exception: |
| return '', None |
| instance_prompt = getattr(card.data, 'instance_prompt', '') |
| image_url = get_image_from_card(card, lora_model_id) |
| return instance_prompt, image_url |
| except Exception as e: |
| raise type(e)(f'failed to load_model_info, due to: {e}') |
|
|
| def update_model_info(self, model_source: str): |
| try: |
| if model_source == 'None': |
| return '', None |
| else: |
| model_info = self.load_model_info(model_source) |
| new_prompt, new_image = model_info[0], model_info[1] |
| return new_prompt, new_image |
| except Exception as e: |
| raise type(e)(f'failed to update_model_info, due to: {e}') |
|
|
|
|
| hf_token = os.getenv('HF_TOKEN') |
| pipe = InferencePipeline(hf_token) |
| app = InferenceUtil(hf_token) |
|
|
|
|
| with gr.Blocks(css=css) as demo: |
| title = gr.HTML( |
| '''<h1>Implicit Style-Content Separation using B-LoRA</h1> |
| <p>This is a demo for our <a href="https://arxiv.org/abs/2403.14572">paper</a>: <b>''Implicit Style-Content Separation using B-LoRA''</b>. |
| <br> |
| Project page and code is available <a href="https://b-lora.github.io/B-LoRA/">here</a>.</p> |
| Select your favorite style and content components from the list. (prefixed with <strong>`B-LoRA-`<strong>) |
| ''', |
| elem_id="title" |
| ) |
| with gr.Row(elem_classes="gr-row"): |
| with gr.Column(): |
| with gr.Group(elem_classes="lora-column"): |
| content_sub_title = gr.HTML('''<h2>Content B-LoRA</h2>''', elem_classes="lora-title") |
| content_checkbox = gr.Checkbox(label='Use Content Only', value=False) |
| content_lora_model_id = gr.Dropdown(label='Model ID', choices=[]) |
| content_prompt = gr.Text(label='Content instance prompt', interactive=False, max_lines=1) |
| content_image = gr.Image(label='Content Image', elem_classes="gr-image") |
| with gr.Column(): |
| with gr.Group(elem_classes="lora-column"): |
| style_sub_title = gr.HTML('''<h2>Style B-LoRA</h2>''', elem_classes="lora-title") |
| style_checkbox = gr.Checkbox(label='Use Style Only', value=False) |
| style_lora_model_id = gr.Dropdown(label='Model ID', choices=[]) |
| style_prompt = gr.Text(label='Style instance prompt', interactive=False, max_lines=1) |
| style_image = gr.Image(label='Style Image', elem_classes="gr-image") |
| with gr.Row(elem_classes="gr-row"): |
| with gr.Column(): |
| with gr.Group(): |
| prompt = gr.Textbox( |
| label='Prompt', |
| max_lines=1, |
| placeholder='Example: "A [c] in [s] style"' |
| ) |
| result = gr.Gallery(label='Result', elem_classes="res-image") |
| with gr.Accordion('Other Parameters', open=False, elem_classes="gr-accordion"): |
| content_alpha = gr.Slider(label='Content B-LoRA alpha', |
| minimum=0, |
| maximum=2, |
| step=0.05, |
| value=1) |
| style_alpha = gr.Slider(label='Style B-LoRA alpha', |
| minimum=0, |
| maximum=2, |
| step=0.05, |
| value=1) |
| seed = gr.Slider(label='Seed', |
| minimum=0, |
| maximum=100000, |
| step=1, |
| value=8888) |
| num_steps = gr.Slider(label='Number of Steps', |
| minimum=0, |
| maximum=100, |
| step=1, |
| value=40) |
| guidance_scale = gr.Slider(label='CFG Scale', |
| minimum=0, |
| maximum=50, |
| step=0.1, |
| value=7.5) |
| num_images_per_prompt = gr.Slider(label='Number of Images per Prompt', |
| minimum=1, |
| maximum=4, |
| step=1, |
| value=2) |
| run_button = gr.Button('Generate') |
| demo.load(demo_init, inputs=[], |
| outputs=[content_lora_model_id, content_prompt, content_image, style_lora_model_id, style_prompt, |
| style_image, prompt], queue=False, show_progress="hidden") |
| content_lora_model_id.change( |
| fn=app.update_model_info, |
| inputs=content_lora_model_id, |
| outputs=[ |
| content_prompt, |
| content_image, |
| ]) |
| style_lora_model_id.change( |
| fn=app.update_model_info, |
| inputs=style_lora_model_id, |
| outputs=[ |
| style_prompt, |
| style_image, |
| ]) |
| style_prompt.change( |
| fn=handle_prompt_change, |
| inputs=[content_prompt, style_prompt], |
| outputs=prompt, |
| ) |
| content_prompt.change( |
| fn=handle_prompt_change, |
| inputs=[content_prompt, style_prompt], |
| outputs=prompt, |
| ) |
| content_checkbox.change(toggle_column, inputs=[content_checkbox], |
| outputs=[style_lora_model_id]) |
| style_checkbox.change(toggle_column, inputs=[style_checkbox], |
| outputs=[content_lora_model_id]) |
| inputs = [ |
| content_lora_model_id, |
| style_lora_model_id, |
| prompt, |
| content_alpha, |
| style_alpha, |
| seed, |
| num_steps, |
| guidance_scale, |
| num_images_per_prompt |
| ] |
| prompt.submit(fn=pipe.run, inputs=inputs, outputs=result) |
| run_button.click(fn=pipe.run, inputs=inputs, outputs=result) |
|
|
| demo.queue(max_size=10).launch(share=False) |
|
|
