Update app.py

app.py

@@ -1,28 +1,34 @@
 import torch
 from diffusers.utils import load_image
 from diffusers import FluxControlNetPipeline, FluxControlNetModel
+import os
 
 # Clear unnecessary memory
 torch.cuda.empty_cache()
 
 # Set the environment variable to handle memory fragmentation
-import os
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
 base_model = 'black-forest-labs/FLUX.1-dev'
 controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
-# Use a smaller precision if possible
+# Use mixed precision with float16
 controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
 pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.float16)
+
+# Offload some parts to CPU
+pipe.enable_model_cpu_offload()
+
 pipe.to("cuda")
 
 control_image_canny = load_image("https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union-alpha/resolve/main/images/canny.jpg")
 controlnet_conditioning_scale = 0.5
 control_mode = 0
 
-# Ensure that image is loaded correctly
+# Reduce the image resolution if needed to fit into memory
 width, height = control_image_canny.size
+width = width // 2
+height = height // 2
 
 prompt = 'A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.'