Update app.py

app.py

@@ -1,249 +1,182 @@
 import gradio as gr
-from time import sleep, time
+from huggingface_hub import login
 from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
-from huggingface_hub import hf_hub_download, CommitScheduler
-from safetensors.torch import load_file
-from share_btn import community_icon_html, loading_icon_html, share_js
-from uuid import uuid4
-from pathlib import Path
-from PIL import Image
 import torch
-import json
-import random
 import copy
-import gc
-import pickle
+import os
 import spaces
+import random
 
-lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
-
-IMAGE_DATASET_DIR = Path("image_dataset") / f"train-{uuid4()}"
-IMAGE_DATASET_DIR.mkdir(parents=True, exist_ok=True)
-IMAGE_JSONL_PATH = IMAGE_DATASET_DIR / "metadata.jsonl"
-
-scheduler = CommitScheduler(
-    repo_id="multimodalart/lora-fusing-preferences",
-    repo_type="dataset",
-    folder_path=IMAGE_DATASET_DIR,
-    path_in_repo=IMAGE_DATASET_DIR.name,
-    every=10
-)
-
-with open(lora_list, "r") as file:
-    data = json.load(file)
-    sdxl_loras = [
-        {
-            "image": item["image"] if item["image"].startswith("https://") else f'https://huggingface.co/spaces/multimodalart/LoraTheExplorer/resolve/main/{item["image"]}',
-            "title": item["title"],
-            "repo": item["repo"],
-            "trigger_word": item["trigger_word"],
-            "weights": item["weights"],
-            "is_compatible": item["is_compatible"],
-            "is_pivotal": item.get("is_pivotal", False),
-            "text_embedding_weights": item.get("text_embedding_weights", None),
-            "is_nc": item.get("is_nc", False)
-        }
-        for item in data
-    ]
-
-state_dicts = {}
-
-for item in sdxl_loras:
-    saved_name = hf_hub_download(item["repo"], item["weights"])
-    
-    if not saved_name.endswith('.safetensors'):
-        state_dict = torch.load(saved_name, map_location=torch.device('cpu'))
-    else:
-        state_dict = load_file(saved_name, device="cpu")
-    
-    state_dicts[item["repo"]] = {
-        "saved_name": saved_name,
-        "state_dict": state_dict
-    }
-
-css = '''
-.gradio-container{max-width: 650px! important}
-#title{text-align:center;}
-#title h1{font-size: 250%}
-.selected_random img{object-fit: cover}
-.selected_random [data-testid="block-label"] span{display: none}
-.plus_column{align-self: center}
-.plus_button{font-size: 235% !important; text-align: center;margin-bottom: 19px}
-#prompt{padding: 0 0 1em 0}
-#prompt input{width: calc(100% - 160px);border-top-right-radius: 0px;border-bottom-right-radius: 0px;}
-#run_button{position: absolute;margin-top: 25.8px;right: 0;margin-right: 0.75em;border-bottom-left-radius: 0px;border-top-left-radius: 0px}
-.random_column{align-self: center; align-items: center;gap: 0.5em !important}
-#share-btn-container{padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; max-width: 13rem; margin-left: auto;margin-top: 0.35em;}
-div#share-btn-container > div {flex-direction: row;background: black;align-items: center}
-#share-btn-container:hover {background-color: #060606}
-#share-btn {all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.5rem !important; padding-bottom: 0.5rem !important;right:0;font-size: 15px;}
-#share-btn * {all: unset}
-#share-btn-container div:nth-child(-n+2){width: auto !important;min-height: 0px !important;}
-#share-btn-container .wrap {display: none !important}
-#share-btn-container.hidden {display: none!important}
-#post_gen_info{margin-top: .5em}
-#thumbs_up_clicked{background:green}
-#thumbs_down_clicked{background:red}
-.title_lora a{color: var(--body-text-color) !important; opacity:0.6}
-#prompt_area .form{border:0}
-#reroll_button{position: absolute;right: 0;flex-grow: 1;min-width: 75px;padding: .1em}
-.pending .min {min-height: auto}
-'''
+hf_token = os.environ.get("HF_TOKEN")
+login(token = hf_token)
 
 original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 
 @spaces.GPU
-def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed=-1):
-  
-  repo_id_1 = shuffled_items[0]['repo']
-  repo_id_2 = shuffled_items[1]['repo']
-  print("Loading state dicts...")
-  start_time = time()  
-  state_dict_1 = copy.deepcopy(state_dicts[repo_id_1]["state_dict"])
-  state_dict_1 = {k: v.to(device="cuda", dtype=torch.float16) for k,v in state_dict_1.items() if torch.is_tensor(v)}
-  state_dict_2 = copy.deepcopy(state_dicts[repo_id_2]["state_dict"])
-  state_dict_2 = {k: v.to(device="cuda", dtype=torch.float16) for k,v in state_dict_2.items() if torch.is_tensor(v)}
-  state_dict_time = time() - start_time
-  print(f"State Dict time: {state_dict_time}")  
-  start_time = time()
-  unet = copy.deepcopy(original_pipe.unet)
-  text_encoder=copy.deepcopy(original_pipe.text_encoder)
-  text_encoder_2=copy.deepcopy(original_pipe.text_encoder_2)
-  pipe = StableDiffusionXLPipeline(vae=original_pipe.vae,
-                                   text_encoder=text_encoder,
-                                   text_encoder_2=text_encoder_2,
-                                   scheduler=original_pipe.scheduler,
-                                   tokenizer=original_pipe.tokenizer,
-                                   tokenizer_2=original_pipe.tokenizer_2,
-                                   unet=unet)
-  pickle_time = time() - start_time
-  print(f"copy time: {pickle_time}")  
-  pipe.to("cuda") 
-  start_time = time()
-  print("Loading LoRA weights...")
-  pipe.load_lora_weights(state_dict_1, low_cpu_mem_usage=True)
-  pipe.fuse_lora(lora_1_scale)
-  pipe.load_lora_weights(state_dict_2, low_cpu_mem_usage=True)
-  pipe.fuse_lora(lora_2_scale)
-  lora_time = time() - start_time
-  print(f"Loaded LoRAs time: {lora_time}") 
-  if negative_prompt == "":
-    negative_prompt = None
-          
-  if(seed < 0):
-      seed = random.randint(0, 2147483647)
-  generator = torch.Generator(device="cuda").manual_seed(seed)
-  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768, generator=generator).images[0]
-  return image, gr.update(visible=True), seed, gr.update(visible=True, interactive=True), gr.update(visible=False), gr.update(visible=True, interactive=True), gr.update(visible=False)
-
-def get_description(item):
-      trigger_word = item["trigger_word"]
-      return f"Trigger: `{trigger_word}`" if trigger_word else "No trigger, applied automatically", trigger_word
-
-def truncate_string(s, max_length=29):
-    return s[:max_length - 3] + "..." if len(s) > max_length else s
-
-def shuffle_images():
-    compatible_items = [item for item in sdxl_loras if item['is_compatible']]
-    random.shuffle(compatible_items)
-    two_shuffled_items = compatible_items[:2]
-    title_1 = gr.update(label=two_shuffled_items[0]['title'], value=two_shuffled_items[0]['image'])
-    title_2 = gr.update(label=two_shuffled_items[1]['title'], value=two_shuffled_items[1]['image'])
-    repo_id_1 = gr.update(value=two_shuffled_items[0]['repo'])
-    repo_id_2 = gr.update(value=two_shuffled_items[1]['repo'])
-    description_1, trigger_word_1 = get_description(two_shuffled_items[0])
-    description_2, trigger_word_2 = get_description(two_shuffled_items[1])
-
-    lora_1_link = f"[{truncate_string(two_shuffled_items[0]['repo'])}](https://huggingface.co/{two_shuffled_items[0]['repo']}) ✨"
-    lora_2_link = f"[{truncate_string(two_shuffled_items[1]['repo'])}](https://huggingface.co/{two_shuffled_items[1]['repo']}) ✨"
-    prompt_description_1 = gr.update(value=description_1, visible=True)
-    prompt_description_2 = gr.update(value=description_2, visible=True)
-    prompt = gr.update(value=f"{trigger_word_1} {trigger_word_2}")
-    scale = gr.update(value=0.7)
-
-    return lora_1_link, title_1, prompt_description_1, repo_id_1, lora_2_link, title_2, prompt_description_2, repo_id_2, prompt, two_shuffled_items, scale, scale
-
-def save_preferences(lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, generated_image, thumbs_direction, seed):
-    image_path = IMAGE_DATASET_DIR / f"{uuid4()}.png"
-    with scheduler.lock:
-        Image.fromarray(generated_image).save(image_path)
-        with IMAGE_JSONL_PATH.open("a") as f:
-            json.dump({"prompt": prompt, "file_name":image_path.name, "lora_1_id": lora_1_id, "lora_1_scale": float(lora_1_scale), "lora_2_id": lora_2_id, "lora_2_scale": float(lora_2_scale), "thumbs_direction": thumbs_direction, "seed": int(seed)}, f)
-            f.write("\n")
+def infer(lora_1_id, lora_1_sfts, lora_2_id, lora_2_sfts, prompt, negative_prompt, lora_1_scale, lora_2_scale, seed):
+
+    unet = copy.deepcopy(original_pipe.unet)
+    text_encoder = copy.deepcopy(original_pipe.text_encoder)
+    text_encoder_2 = copy.deepcopy(original_pipe.text_encoder_2)
+
+    pipe = StableDiffusionXLPipeline(
+        vae = original_pipe.vae,
+        text_encoder = text_encoder,
+        text_encoder_2 = text_encoder_2,
+        scheduler = original_pipe.scheduler,
+        tokenizer = original_pipe.tokenizer,
+        tokenizer_2 = original_pipe.tokenizer_2,
+        unet = unet
+    )
+
+    pipe.to("cuda")
+
+    pipe.load_lora_weights(
+        lora_1_id,
+        weight_name = lora_1_sfts,        
+        low_cpu_mem_usage = True,
+        use_auth_token = True
+    )
+
+    pipe.fuse_lora(lora_1_scale)
+
+    pipe.load_lora_weights(
+        lora_2_id,
+        weight_name = lora_2_sfts,        
+        low_cpu_mem_usage = True,
+        use_auth_token = True
+    )
+
+    pipe.fuse_lora(lora_2_scale)
+
+    if negative_prompt == "" :
+        negative_prompt = None
     
-    return gr.update(visible=False), gr.update(visible=True), gr.update(interactive=False)
-
-def hide_post_gen_info():
-    return gr.update(visible=False)
-
-with gr.Blocks(css=css) as demo:
-  shuffled_items = gr.State()
-  title = gr.HTML(
-        '''<h1>LoRA Roulette 🎰</h1>
-        <p>This random LoRAs are loaded into SDXL, can you find a fun way to combine them? 🎨</p>
-        ''',
-        elem_id="title"
-  )
-  with gr.Column():
-    with gr.Column(min_width=10, scale=16, elem_classes="plus_column"):
+    if seed < 0 :
+        seed = random.randit(0, 423538377342)
+    
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    image = pipe(
+        prompt = prompt,
+        negative_prompt = negative_prompt,
+        num_inference_steps = 25,
+        width = 1024,
+        height = 1024,
+        generator = generator
+    ).images[0]
+    
+    return image, seed
+
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+
+        title = gr.HTML(
+        '''
+        <h1 style="text-align: center;">LoRA Fusion</h1>
+        <p style="text-align: center;">Fuse 2 custom LoRa models</p>
+        '''
+        )
+        
+        # PART 1 • MODELS
         with gr.Row():
-            with gr.Column(min_width=10, scale=4, elem_classes="random_column"):
-              lora_1_link = gr.Markdown(elem_classes="title_lora")
-              lora_1 = gr.Image(interactive=False, height=150, elem_classes="selected_random", elem_id="randomLoRA_1", show_share_button=False, show_download_button=False)
-              lora_1_id = gr.Textbox(visible=False, elem_id="random_lora_1_id")
-              lora_1_prompt = gr.Markdown(visible=False)
-            with gr.Column(min_width=10, scale=1, elem_classes="plus_column"):
-              plus = gr.HTML("+", elem_classes="plus_button")
-            with gr.Column(min_width=10, scale=4, elem_classes="random_column"):
-              lora_2_link = gr.Markdown(elem_classes="title_lora")
-              lora_2 = gr.Image(interactive=False, height=150, elem_classes="selected_random", elem_id="randomLoRA_2", show_share_button=False, show_download_button=False)
-              lora_2_id = gr.Textbox(visible=False, elem_id="random_lora_2_id")
-              lora_2_prompt = gr.Markdown(visible=False)
-            with gr.Column(min_width=10, scale=2, elem_classes="plus_column"):
-               equal = gr.HTML("=", elem_classes="plus_button")
-            shuffle_button = gr.Button("🎲  reroll", elem_id="reroll_button")
-    with gr.Column(min_width=10, scale=14):
-        with gr.Box(elem_id="generate_area"):
-            with gr.Row(elem_id="prompt_area"):
-                prompt = gr.Textbox(label="Your prompt", info="Rearrange the trigger words into a coherent prompt", show_label=False, interactive=True, elem_id="prompt")
-                run_btn = gr.Button("Run", elem_id="run_button")
-            output_image = gr.Image(label="Output", height=460, elem_id="output_image", interactive=False)
-            with gr.Column(visible=False, elem_id="post_gen_info") as post_gen_info:
-                with gr.Row():
-                    with gr.Column(min_width=10):
-                        thumbs_up = gr.Button("👍", elem_id="thumbs_up_unclicked")
-                        thumbs_up_clicked = gr.Button("👍", elem_id="thumbs_up_clicked", interactive=False, visible=False)
-                    with gr.Column(min_width=10):
-                        thumbs_down = gr.Button("👎", elem_id="thumbs_down_unclicked")
-                        thumbs_down_clicked = gr.Button("👎", elem_id="thumbs_down_clicked", interactive=False, visible=False)
-                    with gr.Column(min_width=10):
-                        with gr.Group(elem_id="share-btn-container") as share_group:
-                            community_icon = gr.HTML(community_icon_html)
-                            loading_icon = gr.HTML(loading_icon_html)
-                            share_button = gr.Button("Share to community", elem_id="share-btn")
-                gr.Markdown('<p style="font-size: 90%;opacity: 0.8;">Rating helps finding the most compatible LoRAs 🤗, results are shared annonymously <a href="https://huggingface.co/datasets/multimodalart/lora-fusing-preferences" target="_blank" rel="noopener noreferrer">here</a></p>')
-  with gr.Accordion("Advanced settings", open=False):
-    with gr.Row():
-      lora_1_scale = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=1, step=0.1, value=0.7)
-      lora_2_scale = gr.Slider(label="LoRa 2 Scale", minimum=0, maximum=1, step=0.1, value=0.7)
-    negative_prompt = gr.Textbox(label="Negative prompt")
-    seed = gr.Slider(label="Seed", info="-1 denotes a random seed", minimum=-1, maximum=2147483647, value=-1)
-    last_used_seed = gr.Number(label="Last used seed", info="The seed used in the last generation", minimum=0, maximum=2147483647, value=-1, interactive=False)
-  gr.Markdown("Generate with intent in [LoRA the Explorer](https://huggingface.co/spaces/multimodalart/LoraTheExplorer), but remember: sometimes restrictions flourish creativity 🌸")
+            
+            with gr.Column():
+                
+                lora_1_id = gr.Textbox(
+                    label = "LoRa 1 ID",
+                    placeholder = "username/model_id"
+                )
+
+                lora_1_sfts = gr.Textbox(
+                    label = "Safetensors file",
+                    placeholder = "specific_chosen.safetensors"
+                )
+            
+            with gr.Column():
+                
+                lora_2_id = gr.Textbox(
+                    label = "LoRa 2 ID",
+                    placeholder = "username/model_id"
+                )
+
+                lora_2_sfts = gr.Textbox(
+                    label = "Safetensors file",
+                    placeholder = "specific_chosen.safetensors"
+                )
+
+        # PART 2 • INFERENCE
+        with gr.Row():
+            
+            prompt = gr.Textbox(
+                label = "Your prompt",
+                info = "Use your trigger words into a coherent prompt"
+                placeholde = "e.g: a triggerWordOne portrait in triggerWord2 style"
+            )
+
+            run_btn = gr.Button("Run")
+        
+        output_image = gr.Image(
+            label = "Output"
+        )
+
+        # Advanced Settings
+        with gr.Accordion("Advanced Settings", open=False):
+            
+            with gr.Row():
+                
+                lora_1_scale = gr.Slider(
+                    label = "LoRa 1 scale",
+                    minimum = 0,
+                    maximum = 1,
+                    steps = 0.1,
+                    value = 0.7
+                )
+                
+                lora_2_scale = gr.Slider(
+                    label = "LoRa 2 scale",
+                    minimum = 0,
+                    maximum = 1,
+                    steps = 0.1,
+                    value = 0.7
+                )
+            
+            negative_prompt = gr.Textbox(
+                label = "Negative prompt"
+            )
+
+            seed = gr.Slider(
+                label = "Seed",
+                info = "-1 denotes a random seed",
+                minimum = -1,
+                maximum = 423538377342,
+                value = -1
+            )
+
+            last_used_seed = gr.Number(
+                label = Last used seed,
+                info = "the seed used in the last generation",
+            )
     
-  demo.load(shuffle_images, inputs=[], outputs=[lora_1_link, lora_1, lora_1_prompt, lora_1_id, lora_2_link, lora_2, lora_2_prompt, lora_2_id, prompt, shuffled_items, lora_1_scale, lora_2_scale], queue=False, show_progress="hidden")
-  shuffle_button.click(shuffle_images, outputs=[lora_1_link, lora_1, lora_1_prompt, lora_1_id, lora_2_link, lora_2, lora_2_prompt, lora_2_id, prompt, shuffled_items, lora_1_scale, lora_2_scale], queue=False, show_progress="hidden")
-
-  run_btn.click(hide_post_gen_info, outputs=[post_gen_info], queue=False).then(merge_and_run,
-                inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale, seed],
-                outputs=[output_image, post_gen_info, last_used_seed, thumbs_up, thumbs_up_clicked, thumbs_down, thumbs_down_clicked])
-  prompt.submit(hide_post_gen_info, outputs=[post_gen_info], queue=False).then(merge_and_run,
-                inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale, seed],
-                outputs=[output_image, post_gen_info, last_used_seed, thumbs_up, thumbs_up_clicked, thumbs_down, thumbs_down_clicked])
-  
-  thumbs_up.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("up"), last_used_seed], outputs=[thumbs_up, thumbs_up_clicked, thumbs_down])
-  thumbs_down.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("down"), last_used_seed], outputs=[thumbs_down, thumbs_down_clicked, thumbs_up])
-  share_button.click(None, [], [], _js=share_js)
-
-demo.queue(concurrency_count=2)
-demo.launch()
+    # ACTIONS
+    run_btn.click(
+        fn = infer,
+        inputs = [
+            lora_1_id,
+            lora_1_sfts,
+            lora_2_id,
+            lora_2_sfts,
+            prompt,
+            negative_prompt,
+            lora_1_scale,
+            lora_2_scale,
+            seed
+        ],
+        outputs = [
+            output_image, 
+            last_used_seed
+        ]
+    )
+
+demo.queue(concurrency_count=2).launch()
+