Update app.py

app.py

@@ -25,17 +25,18 @@ base_model = "black-forest-labs/FLUX.1-dev"
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
-pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
-                                                      vae=good_vae,
-                                                      transformer=pipe.transformer,
-                                                      text_encoder=pipe.text_encoder,
-                                                      tokenizer=pipe.tokenizer,
-                                                      text_encoder_2=pipe.text_encoder_2,
-                                                      tokenizer_2=pipe.tokenizer_2,
-                                                      torch_dtype=dtype
-                                                     )
-
-MAX_SEED = 2**32-1
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+    base_model,
+    vae=good_vae,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+    torch_dtype=dtype
+)
+
+MAX_SEED = 2**32 - 1
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
@@ -46,7 +47,7 @@ class calculateDuration:
     def __enter__(self):
         self.start_time = time.time()
         return self
-    
+
     def __exit__(self, exc_type, exc_value, traceback):
         self.end_time = time.time()
         self.elapsed_time = self.end_time - self.start_time
@@ -66,7 +67,7 @@ def update_selection(evt: gr.SelectData, selected_indices, width, height):
             selected_indices.append(selected_index)
         else:
             gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
-            return(
+            return (
                 gr.update(),
                 gr.update(),
                 gr.update(),
@@ -80,19 +81,19 @@ def update_selection(evt: gr.SelectData, selected_indices, width, height):
             )
 
     # Initialize outputs
-    selected_info_1 = ""
-    selected_info_2 = ""
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
     lora_scale_1 = 0.95
     lora_scale_2 = 0.95
     lora_image_1 = None
     lora_image_2 = None
     if len(selected_indices) >= 1:
         lora1 = loras[selected_indices[0]]
-        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
         lora2 = loras[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
 
     # Update prompt placeholder based on last selected LoRA
@@ -128,11 +129,11 @@ def remove_lora_1(selected_indices):
     lora_image_2 = None
     if len(selected_indices) >= 1:
         lora1 = loras[selected_indices[0]]
-        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
         lora2 = loras[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
@@ -149,11 +150,11 @@ def remove_lora_2(selected_indices):
     lora_image_2 = None
     if len(selected_indices) >= 1:
         lora1 = loras[selected_indices[0]]
-        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
         lora2 = loras[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
@@ -163,8 +164,8 @@ def randomize_loras(selected_indices):
     selected_indices = random.sample(range(len(loras)), 2)
     lora1 = loras[selected_indices[0]]
     lora2 = loras[selected_indices[1]]
-    selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
-    selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+    selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+    selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
     lora_scale_1 = 0.95
     lora_scale_2 = 0.95
     lora_image_1 = lora1['image']
@@ -173,7 +174,7 @@ def randomize_loras(selected_indices):
 
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
-    print("Entrou aqui!")
+    print("Generating image...")
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
     with calculateDuration("Generating image"):
@@ -208,8 +209,8 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
         joint_attention_kwargs={"scale": 1.0},
         output_type="pil",
     ).images[0]
-    return final_image 
-    
+    return final_image
+
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, randomize_seed, seed, width, height, progress=gr.Progress(track_tqdm=True)):
     if not selected_indices:
         raise gr.Error("You must select at least one LoRA before proceeding.")
@@ -235,27 +236,29 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
 
     # Load LoRA weights with respective scales
     lora_names = []
+    lora_weights = []
     with calculateDuration("Loading LoRA weights"):
         for idx, lora in enumerate(selected_loras):
             lora_name = f"lora_{idx}"
             lora_names.append(lora_name)
+            lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2)
             lora_path = lora['repo']
-            scale = lora_scale_1 if idx == 0 else lora_scale_2
+            weight_name = lora.get("weights")
             if image_input is not None:
-                if "weights" in lora:
-                    pipe_i2i.load_lora_weights(lora_path, weight_name=lora["weights"], low_cpu_mem_usage=True, adapter_name=lora_name)
+                if weight_name:
+                    pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name)
                 else:
                     pipe_i2i.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name)
             else:
-                if "weights" in lora:
-                    pipe.load_lora_weights(lora_path, weight_name=lora["weights"], low_cpu_mem_usage=True, adapter_name=lora_name)
+                if weight_name:
+                    pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name=lora_name)
                 else:
                     pipe.load_lora_weights(lora_path, low_cpu_mem_usage=True, adapter_name=lora_name)
-        print(lora_names)
+        print("Loaded LoRAs:", lora_names)
         if image_input is not None:
-            pipe_i2i.set_adapters(lora_names, adapter_weights=[lora_scale_1, lora_scale_2])
+            pipe_i2i.set_adapters(lora_names, adapter_weights=lora_weights)
         else:
-            pipe.set_adapters(lora_names, adapter_weights=[lora_scale_1, lora_scale_2])
+            pipe.set_adapters(lora_names, adapter_weights=lora_weights)
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
@@ -271,7 +274,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         final_image = None
         step_counter = 0
         for image in image_generator:
-            step_counter+=1
+            step_counter += 1
             final_image = image
             progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
             yield image, seed, gr.update(value=progress_bar, visible=True)
@@ -282,7 +285,7 @@ def get_huggingface_safetensors(link):
     if len(split_link) == 2:
         model_card = ModelCard.load(link)
         base_model = model_card.data.get("base_model")
-        print(base_model)
+        print(f"Base model: {base_model}")
         if base_model not in ["black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"]:
             raise Exception("Not a FLUX LoRA!")
         image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
@@ -304,13 +307,26 @@ def get_huggingface_safetensors(link):
         if not safetensors_name:
             raise Exception("No *.safetensors file found in the repository")
         return split_link[1], link, safetensors_name, trigger_word, image_url
+    else:
+        raise Exception("Invalid Hugging Face repository link")
 
 def check_custom_model(link):
-    if link.startswith("https://"):
-        if link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co"):
+    if link.endswith(".safetensors"):
+        # Treat as direct link to the LoRA weights
+        title = os.path.basename(link)
+        repo = link
+        path = None  # No specific weight name
+        trigger_word = ""
+        image_url = None
+        return title, repo, path, trigger_word, image_url
+    elif link.startswith("https://"):
+        if "huggingface.co" in link:
             link_split = link.split("huggingface.co/")
             return get_huggingface_safetensors(link_split[1])
-    else: 
+        else:
+            raise Exception("Unsupported URL")
+    else:
+        # Assume it's a Hugging Face model path
         return get_huggingface_safetensors(link)
 
 def add_custom_lora(custom_lora, selected_indices):
@@ -319,18 +335,6 @@ def add_custom_lora(custom_lora, selected_indices):
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
             print(f"Loaded custom LoRA: {repo}")
-            card = f'''
-            <div class="custom_lora_card">
-              <span>Loaded custom LoRA:</span>
-              <div class="card_internal">
-                <img src="{image}" />
-                <div>
-                    <h3>{title}</h3>
-                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
-                </div>
-              </div>
-            </div>
-            '''
             existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
             if existing_item_index is None:
                 new_item = {
@@ -340,7 +344,7 @@ def add_custom_lora(custom_lora, selected_indices):
                     "weights": path,
                     "trigger_word": trigger_word
                 }
-                print(new_item)
+                print(f"New LoRA: {new_item}")
                 existing_item_index = len(loras)
                 loras.append(new_item)
 
@@ -349,32 +353,42 @@ def add_custom_lora(custom_lora, selected_indices):
             # Update selected_indices if there's room
             if len(selected_indices) < 2:
                 selected_indices.append(existing_item_index)
-                selected_info_1 = ""
-                selected_info_2 = ""
-                lora_scale_1 = 0.95
-                lora_scale_2 = 0.95
-                lora_image_1 = None
-                lora_image_2 = None
-                if len(selected_indices) >= 1:
-                    lora1 = loras[selected_indices[0]]
-                    selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
-                    lora_image_1 = lora1['image']
-                if len(selected_indices) >= 2:
-                    lora2 = loras[selected_indices[1]]
-                    selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
-                    lora_image_2 = lora2['image']
-                return (gr.update(visible=True, value=card), gr.update(visible=True), gr.update(value=gallery_items),
-                        selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2)
             else:
-                return (gr.update(visible=True, value=card), gr.update(visible=True), gr.update(value=gallery_items),
-                        gr.NoChange(), gr.NoChange(), selected_indices, gr.NoChange(), gr.NoChange(), gr.NoChange(), gr.NoChange())
+                gr.Warning("You can select up to 2 LoRAs, remove one to select a new one.")
+
+            # Update selected_info and images
+            selected_info_1 = "Select a LoRA 1"
+            selected_info_2 = "Select a LoRA 2"
+            lora_scale_1 = 0.95
+            lora_scale_2 = 0.95
+            lora_image_1 = None
+            lora_image_2 = None
+            if len(selected_indices) >= 1:
+                lora1 = loras[selected_indices[0]]
+                selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
+                lora_image_1 = lora1['image']
+            if len(selected_indices) >= 2:
+                lora2 = loras[selected_indices[1]]
+                selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
+                lora_image_2 = lora2['image']
+            return (
+                gr.update(value=gallery_items),
+                selected_info_1,
+                selected_info_2,
+                selected_indices,
+                lora_scale_1,
+                lora_scale_2,
+                lora_image_1,
+                lora_image_2
+            )
         except Exception as e:
             print(e)
-            return gr.update(visible=True, value=str(e)), gr.update(visible=True), gr.NoChange(), gr.NoChange(), gr.NoChange(), selected_indices, gr.NoChange(), gr.NoChange(), gr.NoChange(), gr.NoChange()
+            gr.Error(str(e))
+            return gr.NoChange(), gr.NoChange(), gr.NoChange(), selected_indices, gr.NoChange(), gr.NoChange(), gr.NoChange(), gr.NoChange()
     else:
-        return gr.update(visible=False), gr.update(visible=False), gr.NoChange(), gr.NoChange(), gr.NoChange(), selected_indices, gr.NoChange(), gr.NoChange(), gr.NoChange(), gr.NoChange()
+        return gr.NoChange(), gr.NoChange(), gr.NoChange(), selected_indices, gr.NoChange(), gr.NoChange(), gr.NoChange(), gr.NoChange()
 
-def remove_custom_lora(custom_lora_info, custom_lora_button, selected_indices):
+def remove_custom_lora(selected_indices):
     global loras
     if loras:
         custom_lora_repo = loras[-1]['repo']
@@ -387,21 +401,30 @@ def remove_custom_lora(custom_lora_info, custom_lora_button, selected_indices):
     # Update gallery
     gallery_items = [(item["image"], item["title"]) for item in loras]
     # Update selected_info and images
-    selected_info_1 = ""
-    selected_info_2 = ""
+    selected_info_1 = "Select a LoRA 1"
+    selected_info_2 = "Select a LoRA 2"
     lora_scale_1 = 0.95
     lora_scale_2 = 0.95
     lora_image_1 = None
     lora_image_2 = None
     if len(selected_indices) >= 1:
         lora1 = loras[selected_indices[0]]
-        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}]({lora1['repo']}) ✨"
         lora_image_1 = lora1['image']
     if len(selected_indices) >= 2:
         lora2 = loras[selected_indices[1]]
-        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}]({lora2['repo']}) ✨"
         lora_image_2 = lora2['image']
-    return gr.update(visible=False), gr.update(visible=False), gr.update(value=gallery_items), selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
+    return (
+        gr.update(value=gallery_items),
+        selected_info_1,
+        selected_info_2,
+        selected_indices,
+        lora_scale_1,
+        lora_scale_2,
+        lora_image_1,
+        lora_image_2
+    )
 
 run_lora.zerogpu = True
 
@@ -410,6 +433,7 @@ css = '''
 #title{text-align: center}
 #title h1{font-size: 3em; display:inline-flex; align-items:center}
 #title img{width: 100px; margin-right: 0.5em}
+#gallery{height: 260px}
 #gallery .grid-wrap{height: 5vh}
 #lora_list{background: var(--block-background-fill);padding: 0 1em .3em; font-size: 90%}
 .custom_lora_card{margin-bottom: 1em}
@@ -460,6 +484,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
                 remove_button_2 = gr.Button("Remove", size="sm")
     with gr.Row():
         with gr.Column():
+            with gr.Group():
+                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path or *.safetensors public URL", placeholder="multimodalart/vintage-ads-flux")
+                add_custom_lora_button = gr.Button("Add Custom LoRA")
+                remove_custom_lora_button = gr.Button("Remove Custom LoRA")
+                gr.Markdown("[Check the list of FLUX LoRAs](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
             gallery = gr.Gallery(
                 [(item["image"], item["title"]) for item in loras],
                 label="LoRA Gallery",
@@ -467,11 +496,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
                 columns=5,
                 elem_id="gallery"
             )
-            with gr.Group():
-                custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path", placeholder="multimodalart/vintage-ads-flux")
-                gr.Markdown("[Check the list of FLUX LoRAs](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.1-dev)", elem_id="lora_list")
-            custom_lora_info = gr.HTML(visible=False)
-            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress", visible=False)
             result = gr.Image(label="Generated Image")
@@ -484,15 +508,15 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
                 with gr.Row():
                     cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
-                
+
                 with gr.Row():
                     width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
                     height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
-                
+
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-    
+
     gallery.select(
         update_selection,
         inputs=[selected_indices, width, height],
@@ -513,15 +537,15 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 3600)) as app:
         inputs=[selected_indices],
         outputs=[selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
-    custom_lora.change(
+    add_custom_lora_button.click(
         add_custom_lora,
         inputs=[custom_lora, selected_indices],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        outputs=[gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
-    custom_lora_button.click(
+    remove_custom_lora_button.click(
         remove_custom_lora,
-        inputs=[custom_lora_info, custom_lora_button, selected_indices],
-        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
+        inputs=[selected_indices],
+        outputs=[gallery, selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2]
     )
     gr.on(
         triggers=[generate_button.click, prompt.submit],