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],