Update app.py

app.py

@@ -4,14 +4,19 @@ from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiCont
 import gradio as gr
 import spaces
 
+# Ensure that you're using the appropriate data type for your GPU
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
 base_model = 'black-forest-labs/FLUX.1-dev'
 controlnet_model_union = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
-controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch.bfloat16)
-controlnet = FluxMultiControlNetModel([controlnet_union])  # we always recommend loading via FluxMultiControlNetModel
+controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch_dtype)
+controlnet = FluxMultiControlNetModel([controlnet_union])
+
+pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch_dtype)
 
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
-pipe.to("cuda")
+# If you encounter issues with CUDA, you can run this on the CPU for debugging
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
 control_modes = [
     "canny",
@@ -46,17 +51,21 @@ def generate_image(prompt, control_image_depth, control_mode_depth_index, use_de
     width, height = control_image_canny.shape[:2]
     adjusted_width, adjusted_height = adjust_dimensions(width, height)
 
-    image = pipe(
-        prompt,
-        control_image=control_images,
-        control_mode=control_modes,
-        width=adjusted_width,
-        height=adjusted_height,
-        controlnet_conditioning_scale=conditioning_scales,
-        num_inference_steps=24,
-        guidance_scale=3.5,
-        generator=torch.manual_seed(42),
-    ).images[0]
+    try:
+        image = pipe(
+            prompt,
+            control_image=control_images,
+            control_mode=control_modes,
+            width=adjusted_width,
+            height=adjusted_height,
+            controlnet_conditioning_scale=conditioning_scales,
+            num_inference_steps=24,
+            guidance_scale=3.5,
+            generator=torch.manual_seed(42),
+        ).images[0]
+    except RuntimeError as e:
+        torch.cuda.empty_cache()
+        raise e
 
     return image
 
@@ -75,4 +84,4 @@ iface = gr.Interface(
     description="Generate an image using FluxControlNet with depth and canny control images.",
 )
 
-iface.launch(share=True)
+iface.launch(share=True)