Update app.py

app.py

@@ -18,6 +18,11 @@ import boto3
 from io import BytesIO
 from datetime import datetime
 
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
+
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
@@ -27,13 +32,9 @@ login(token=HF_TOKEN)
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
-
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
-
 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)
-
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -86,7 +87,8 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
             width=width,
             height=height,
             generator=generator,
-            joint_attention_kwargs={"scale": lora_scale}
+            joint_attention_kwargs={"scale": lora_scale},
+            max_sequence_length=256
         ).images[0]
         
     progress(99, "Generate success!")
@@ -95,8 +97,8 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
 
 def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     
-    with calculateDuration("Unloading LoRA"):
-        pipe.unload_lora_weights()
+    # with calculateDuration("Unloading LoRA"):
+    #     pipe.unload_lora_weights()
         
     # Load LoRA weights
     if lora_repo and lora_name: