Update app.py

app.py

@@ -13,7 +13,6 @@ import torch
 from PIL import Image
 import gradio as gr
 
-
 from diffusers import (
     DiffusionPipeline,
     AutoencoderTiny,
@@ -26,13 +25,21 @@ from huggingface_hub import (
     hf_hub_download,
     HfFileSystem,
     ModelCard,
-    snapshot_download)
+    snapshot_download,
+    login  # imported for one-time authentication
+)
 
 from diffusers.utils import load_image
 
 import spaces
 
-
+# -------------------------------
+# Authenticate with Hugging Face once
+# -------------------------------
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN:
+    login(HF_TOKEN)
+    print("Authenticated with Hugging Face.")
 
 def calculate_shift(
     image_seq_len,
@@ -2089,24 +2096,25 @@ loras = [
 ]
 
 #--------------------------------------------------Model Initialization-----------------------------------------------------------------------------------------#
-
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
-#TAEF1 is very tiny autoencoder which uses the same "latent API" as FLUX.1's VAE. FLUX.1 is useful for real-time previewing of the FLUX.1 generation process.#
-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
-                                                     )
+# TAEF1 is a very tiny autoencoder using the same "latent API" as FLUX.1's VAE.
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype, use_auth_token=HF_TOKEN).to(device)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype, use_auth_token=HF_TOKEN).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1, use_auth_token=HF_TOKEN).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,
+    use_auth_token=HF_TOKEN
+)
 
 MAX_SEED = 2**32-1
 
@@ -2210,7 +2218,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         pipe.unload_lora_weights()
         pipe_i2i.unload_lora_weights()
         
-    #LoRA weights flow
+    # LoRA weights flow
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
         pipe_to_use = pipe_i2i if image_input is not None else pipe
         weight_name = selected_lora.get("weights", None)
@@ -2235,7 +2243,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)
@@ -2243,41 +2251,37 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
         
 def get_huggingface_safetensors(link):
-  split_link = link.split("/")
-  if(len(split_link) == 2):
-            model_card = ModelCard.load(link)
-            base_model = model_card.data.get("base_model")
-            print(base_model)
+    split_link = link.split("/")
+    if len(split_link) == 2:
+        model_card = ModelCard.load(link)
+        base_model = model_card.data.get("base_model")
+        print(base_model)
       
-            #Allows Both
-            if((base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell")):
-                raise Exception("Flux LoRA Not Found!")
+        # Allows Both
+        if (base_model != "black-forest-labs/FLUX.1-dev") and (base_model != "black-forest-labs/FLUX.1-schnell"):
+            raise Exception("Flux LoRA Not Found!")
                 
-            # Only allow "black-forest-labs/FLUX.1-dev"
-            #if base_model != "black-forest-labs/FLUX.1-dev":
-                #raise Exception("Only FLUX.1-dev is supported, other LoRA models are not allowed!")
-                
-            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-            trigger_word = model_card.data.get("instance_prompt", "")
-            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-            fs = HfFileSystem()
-            try:
-                list_of_files = fs.ls(link, detail=False)
-                for file in list_of_files:
-                    if(file.endswith(".safetensors")):
-                        safetensors_name = file.split("/")[-1]
-                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                      image_elements = file.split("/")
-                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-            except Exception as e:
-              print(e)
-              gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-              raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            return split_link[1], link, safetensors_name, trigger_word, image_url
+        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+        trigger_word = model_card.data.get("instance_prompt", "")
+        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+        fs = HfFileSystem()
+        try:
+            list_of_files = fs.ls(link, detail=False)
+            for file in list_of_files:
+                if file.endswith(".safetensors"):
+                    safetensors_name = file.split("/")[-1]
+                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                    image_elements = file.split("/")
+                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+        except Exception as e:
+            print(e)
+            gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+        return split_link[1], link, safetensors_name, trigger_word, image_url
 
 def check_custom_model(link):
-    if(link.startswith("https://")):
-        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
+    if link.startswith("https://"):
+        if link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co"):
             link_split = link.split("huggingface.co/")
             return get_huggingface_safetensors(link_split[1])
     else: 
@@ -2285,7 +2289,7 @@ def check_custom_model(link):
 
 def add_custom_lora(custom_lora):
     global loras
-    if(custom_lora):
+    if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
             print(f"Loaded custom LoRA: {repo}")
@@ -2302,7 +2306,7 @@ def add_custom_lora(custom_lora):
             </div>
             '''
             existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
-            if(not existing_item_index):
+            if existing_item_index is None:
                 new_item = {
                     "image": image,
                     "title": title,
@@ -2316,8 +2320,8 @@ def add_custom_lora(custom_lora):
         
             return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
         except Exception as e:
-            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-FLUX LoRA")
-            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-FLUX LoRA"), gr.update(visible=False), gr.update(), "", None, ""
+            gr.Warning("Invalid LoRA: either you entered an invalid link, or a non-FLUX LoRA")
+            return gr.update(visible=True, value="Invalid LoRA: either you entered an invalid link, a non-FLUX LoRA"), gr.update(visible=False), gr.update(), "", None, ""
     else:
         return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""
 
@@ -2371,7 +2375,7 @@ with gr.Blocks(theme="prithivMLmods/Minecraft-Theme", css=css, delete_cache=(60,
             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)
+            progress_bar = gr.Markdown(elem_id="progress", visible=False)
             result = gr.Image(label="Generated Image")
 
     with gr.Row():