updates

app.py

@@ -1,9 +1,17 @@
-# app.py
-from flux_app.frontend import Frontend  # Import Frontend from the package
+#app.py
+import os
+from flux_app.frontend import Frontend
 from flux_app.backend import ModelManager
+
 if __name__ == "__main__":
-    model_manager = ModelManager()  # Initialize models (backend)
-    frontend = Frontend(model_manager)  # Create frontend, passing the model manager
-    app = frontend.create_ui()  # Create the Gradio app
-    app.queue()  # Important for handling concurrency
-    app.launch()  # Launch the app
+    # Get the Hugging Face token from an environment variable
+    hf_token = os.environ.get("HF_TOKEN")
+
+    if not hf_token:
+        raise ValueError("Hugging Face token (HF_TOKEN) not found in environment variables.  Please set it.")
+    model_manager = ModelManager(hf_token=hf_token)
+
+    frontend = Frontend(model_manager)
+    app = frontend.create_ui()
+    app.queue()
+    app.launch()

flux_app/backend.py

@@ -10,21 +10,28 @@ from flux_app.config import DTYPE, DEVICE, BASE_MODEL, TAEF1_MODEL, MAX_SEED  #
 from flux_app.utilities import calculate_shift, retrieve_timesteps, load_image_from_path, calculateDuration # Absolute import
 from flux_app.lora_handling import  flux_pipe_call_that_returns_an_iterable_of_images # Absolute import
 import time
+from huggingface_hub import login
 
 class ModelManager:
-    def __init__(self):
+    def __init__(self, hf_token=None):
         self.pipe = None
         self.pipe_i2i = None
         self.good_vae = None
         self.taef1 = None
+
+        if hf_token:
+            login(token=hf_token)  # Log in with the provided token
+        #else:  # Optional:  You could add a fallback to interactive login
+        #   login()
+
         self.initialize_models()
 
 
     def initialize_models(self):
         """Initializes the diffusion pipelines and autoencoders."""
-        self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE).to(DEVICE)
-        self.good_vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE).to(DEVICE)
-        self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1).to(DEVICE)
+        self.taef1 = AutoencoderTiny.from_pretrained(TAEF1_MODEL, torch_dtype=DTYPE, token=True).to(DEVICE)
+        self.good_vae = AutoencoderKL.from_pretrained(BASE_MODEL, subfolder="vae", torch_dtype=DTYPE, token=True).to(DEVICE)
+        self.pipe = DiffusionPipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE, vae=self.taef1, token=True).to(DEVICE)
         self.pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
             BASE_MODEL,
             vae=self.good_vae,
@@ -34,6 +41,7 @@ class ModelManager:
             text_encoder_2=self.pipe.text_encoder_2,
             tokenizer_2=self.pipe.tokenizer_2,
             torch_dtype=DTYPE,
+            token=True
         )
         self.pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(self.pipe)
 
@@ -75,4 +83,4 @@ class ModelManager:
                 joint_attention_kwargs={"scale": lora_scale},
                 output_type="pil",
             ).images[0]
-            return final_image  
+            return final_image