bugfix

app.py

@@ -163,12 +163,16 @@ def run_flux(
     progress
 ) -> Image.Image:
     print("Running FLUX...")
+    clear_cuda_cache()
+    # pipe.enable_vae_slicing()
+    # pipe.enable_vae_tiling()
+    pipe.enable_model_cpu_offload()
+
     if lora_path and lora_weights:
         with calculateDuration("load lora"):
             print("start to load lora", lora_path, lora_weights)
-            pipe.unload_lora_weights()
             pipe.load_lora_weights(lora_path, weight_name=lora_weights)
-
+    
     width, height = resolution_wh
     if randomize_seed_checkbox:
         seed_slicer = random.randint(0, MAX_SEED)
@@ -176,25 +180,30 @@ def run_flux(
 
     with calculateDuration("run pipe"):
         print("start to run pipe", prompt)
-        generated_image = pipe(
-            prompt=prompt,
-            image=image,
-            mask_image=mask,
-            control_image=control_image,
-            control_mode=control_mode,
-            controlnet_conditioning_scale=0.85,
-            width=width,
-            height=height,
-            strength=strength_slider,
-            generator=generator,
-            num_inference_steps=num_inference_steps_slider,
-            # max_sequence_length=256,
-            joint_attention_kwargs={"scale": lora_scale}
-        ).images[0]
+       
+        with torch.inference_mode():
+            generated_image = pipe(
+                prompt=prompt,
+                image=image,
+                mask_image=mask,
+                control_image=control_image,
+                control_mode=control_mode,
+                controlnet_conditioning_scale=0.85,
+                width=width,
+                height=height,
+                strength=strength_slider,
+                generator=generator,
+                num_inference_steps=num_inference_steps_slider,
+                # max_sequence_length=256,
+                cross_attention_kwargs={"scale":0.5},
+                joint_attention_kwargs={"scale": lora_scale}
+            ).images[0]
     progress(99, "Generate image success!")
     return generated_image
 
 @spaces.GPU(duration=120)
+@torch.inference_mode()
+
 def process(
     image_url: str,
     mask_url: str,
@@ -315,7 +324,7 @@ def process(
         result["message"] = "generate image failed"
         print(e)
         generated_image = None
-
+    clear_cuda_cache()
     print("run flux finish")
     if generated_image:
         if upload_to_r2:

requirements.txt

@@ -1,6 +1,7 @@
 accelerate
 torch==2.0.1
-transformers
+torchvision==0.15.2
+transformers==4.43.3
 tqdm
 einops
 spaces
@@ -13,4 +14,8 @@ peft
 controlnet-aux
 mediapipe
 kornia
-xformers
+xformers
+einops
+onnxruntime-gpu
+omegaconf
+scipy