Upload 3 files

app_inference.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python
+
+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 = [
+    '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-multi-dog2',
+]
+css = """
+body {
+    font-size: 30px;
+}
+.gr-image {
+    width: 512px;
+    height: 512px;
+    object-fit: contain;
+    margin: auto;
+}
+
+.lora-column {
+    display: flex;
+    flex-direction: column;
+    align-items: center; /* Center align content vertically in columns */
+    justify-content: center; /* Center content horizontally in columns */
+}
+.gr-row {
+    align-items: center;
+    justify-content: center;
+    margin-top: 5px;
+}
+"""
+
+
+def get_choices(hf_token):
+    api = HfApi(token=hf_token)
+    choices = [
+        info.modelId for info in api.list_models(author='lora-library')
+    ]
+    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}')
+
+
+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}')
+
+
+def create_inference_demo(pipe,  #: InferencePipeline,
+                          hf_token: str | None = None) -> gr.Blocks:
+    with gr.Blocks(css=css) as demo:
+        with gr.Row(elem_classes="gr-row"):
+            with gr.Column():
+                with gr.Group(elem_classes="lora-column"):
+                    gr.Markdown('## Content B-LoRA')
+                    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"):
+                    gr.Markdown('## Style B-LoRA')
+                    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.Image(label='Result')
+                    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=50)
+                        guidance_scale = gr.Slider(label='CFG Scale',
+                                                   minimum=0,
+                                                   maximum=50,
+                                                   step=0.1,
+                                                   value=7.5)
+
+                    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=lambda content_blora_prompt,
+                      style_blora_prompt: f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' if style_blora_prompt else content_blora_prompt,
+            inputs=[content_prompt, style_prompt],
+            outputs=prompt,
+        )
+        content_prompt.change(
+            fn=lambda content_blora_prompt,
+                      style_blora_prompt: f'{content_blora_prompt} in {style_blora_prompt[0].lower() + style_blora_prompt[1:]} style' if content_blora_prompt else style_blora_prompt,
+            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,
+        ]
+        prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
+        run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
+    return demo
+
+
+if __name__ == '__main__':
+    hf_token = os.getenv('HF_TOKEN')
+    pipe = InferencePipeline(hf_token)
+    app = InferenceUtil(hf_token)
+    demo = create_inference_demo(pipe, hf_token)
+    demo.queue(max_size=10).launch(share=False)

blora_utils.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+BLOCKS = {
+    'content': ['unet.up_blocks.0.attentions.0'],
+    'style': ['unet.up_blocks.0.attentions.1'],
+}
+
+
+def is_belong_to_blocks(key, blocks):
+    try:
+        for g in blocks:
+            if g in key:
+                return True
+        return False
+    except Exception as e:
+        raise type(e)(f'failed to is_belong_to_block, due to: {e}')
+
+
+def filter_lora(state_dict, blocks_):
+    try:
+        return {k: v for k, v in state_dict.items() if is_belong_to_blocks(k, blocks_)}
+    except Exception as e:
+        raise type(e)(f'failed to filter_lora, due to: {e}')
+
+
+def scale_lora(state_dict, alpha):
+    try:
+        return {k: v * alpha for k, v in state_dict.items()}
+    except Exception as e:
+        raise type(e)(f'failed to scale_lora, due to: {e}')
+
+
+def get_target_modules(unet, blocks=None):
+    try:
+        if not blocks:
+            blocks = [('.').join(blk.split('.')[1:]) for blk in BLOCKS['content'] + BLOCKS['style']]
+
+        attns = [attn_processor_name.rsplit('.', 1)[0] for attn_processor_name, _ in unet.attn_processors.items() if
+                 is_belong_to_blocks(attn_processor_name, blocks)]
+
+        target_modules = [f'{attn}.{mat}' for mat in ["to_k", "to_q", "to_v", "to_out.0"] for attn in attns]
+        return target_modules
+    except Exception as e:
+        raise type(e)(f'failed to get_target_modules, due to: {e}')
+
+

inf.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import gc
+import pathlib
+
+import gradio as gr
+import PIL.Image
+import torch
+from diffusers import StableDiffusionXLPipeline
+from huggingface_hub import ModelCard
+
+from blora_utils import BLOCKS, filter_lora, scale_lora
+
+
+class InferencePipeline:
+    def __init__(self, hf_token: str | None = None):
+        self.hf_token = hf_token
+        self.base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        if self.device.type == 'cpu':
+            self.pipe = StableDiffusionXLPipeline.from_pretrained(
+                self.base_model_id, use_auth_token=self.hf_token, cache_dir='./cache')
+        else:
+            self.pipe = StableDiffusionXLPipeline.from_pretrained(
+                self.base_model_id,
+                torch_dtype=torch.float16,
+                use_auth_token=self.hf_token)
+            self.pipe = self.pipe.to(self.device)
+        self.content_lora_model_id = None
+        self.style_lora_model_id = None
+
+    def clear(self) -> None:
+        self.content_lora_model_id = None
+        self.style_lora_model_id = None
+        del self.pipe
+        self.pipe = None
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    def load_b_lora_to_unet(self, content_lora_model_id: str, style_lora_model_id: str, content_alpha: float,
+                            style_alpha: float) -> None:
+        try:
+            # Get Content B-LoRA SD
+            if content_lora_model_id:
+                content_B_LoRA_sd, _ = self.pipe.lora_state_dict(content_lora_model_id, use_auth_token=self.hf_token)
+                content_B_LoRA = filter_lora(content_B_LoRA_sd, BLOCKS['content'])
+                content_B_LoRA = scale_lora(content_B_LoRA, content_alpha)
+            else:
+                content_B_LoRA = {}
+
+            # Get Style B-LoRA SD
+            if style_lora_model_id:
+                style_B_LoRA_sd, _ = self.pipe.lora_state_dict(style_lora_model_id, use_auth_token=self.hf_token)
+                style_B_LoRA = filter_lora(style_B_LoRA_sd, BLOCKS['style'])
+                style_B_LoRA = scale_lora(style_B_LoRA, style_alpha)
+            else:
+                style_B_LoRA = {}
+
+            # Merge B-LoRAs SD
+            res_lora = {**content_B_LoRA, **style_B_LoRA}
+
+            # Load
+            self.pipe.load_lora_into_unet(res_lora, None, self.pipe.unet)
+        except Exception as e:
+            raise type(e)(f'failed to load_b_lora_to_unet, due to: {e}')
+
+    @staticmethod
+    def check_if_model_is_local(lora_model_id: str) -> bool:
+        return pathlib.Path(lora_model_id).exists()
+
+    @staticmethod
+    def get_model_card(model_id: str,
+                       hf_token: str | None = None) -> ModelCard:
+        if InferencePipeline.check_if_model_is_local(model_id):
+            card_path = (pathlib.Path(model_id) / 'README.md').as_posix()
+        else:
+            card_path = model_id
+        return ModelCard.load(card_path, token=hf_token)
+
+    @staticmethod
+    def get_base_model_info(lora_model_id: str,
+                            hf_token: str | None = None) -> str:
+        card = InferencePipeline.get_model_card(lora_model_id, hf_token)
+        return card.data.base_model
+
+    def load_pipe(self, content_lora_model_id: str, style_lora_model_id: str, content_alpha: float,
+                  style_alpha: float) -> None:
+        if content_lora_model_id == self.content_lora_model_id and style_lora_model_id == self.style_lora_model_id:
+            return
+        self.pipe.unload_lora_weights()
+
+        self.load_b_lora_to_unet(content_lora_model_id, style_lora_model_id, content_alpha, style_alpha)
+
+        self.content_lora_model_id = content_lora_model_id
+        self.style_lora_model_id = style_lora_model_id
+
+    def run(
+            self,
+            content_lora_model_id: str,
+            style_lora_model_id: str,
+            prompt: str,
+            content_alpha: float,
+            style_alpha: float,
+            seed: int,
+            n_steps: int,
+            guidance_scale: float,
+    ) -> PIL.Image.Image:
+        if not torch.cuda.is_available():
+            raise gr.Error('CUDA is not available.')
+
+        self.load_pipe(content_lora_model_id, style_lora_model_id, content_alpha, style_alpha)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        out = self.pipe(
+            prompt,
+            num_inference_steps=n_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )  # type: ignore
+        return out.images[0]