Update app.py

app.py

@@ -2,24 +2,33 @@ import torch
 from diffusers.utils import load_image
 from diffusers import FluxControlNetPipeline, FluxControlNetModel
 
+# 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'
 
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
+# Use a smaller precision if possible
+controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
+pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.float16)
 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
 
-width, height = control_image.size
+# Ensure that image is loaded correctly
+width, height = control_image_canny.size
 
 prompt = 'A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.'
 
 image = pipe(
     prompt, 
-    control_image=control_image,
+    control_image=control_image_canny,
     control_mode=control_mode,
     width=width,
     height=height,
@@ -28,3 +37,6 @@ image = pipe(
     guidance_scale=3.5,
 ).images[0]
 image.save("image.jpg")
+
+# Empty cache after the operation to free up memory
+torch.cuda.empty_cache()