Update app.py

app.py

@@ -5,6 +5,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 from einops import rearrange
+import io
 import requests
 import spaces
 from huggingface_hub import login
@@ -12,20 +13,17 @@ from gradio_imageslider import ImageSlider
 from diffusers.utils import load_image
 from diffusers import FluxControlNetPipeline, FluxControlNetModel
 
-# Source: https://github.com/XLabs-AI/x-flux.git
-name = "flux-dev"
-
-# Load the model on CPU
+# Device settings: CPU for loading, GPU for inference
 device_cpu = torch.device("cpu")
-device_gpu = torch.device("cuda")
+device_gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Model identifiers
 base_model = 'black-forest-labs/FLUX.1-dev'
 controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
 # Load the ControlNet model and pipeline on CPU
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16, device_map="cpu")
-pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
-pipe.to(device_cpu)  # Keep on CPU initially
+controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16).to(device_cpu)
+pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16).to(device_cpu)
 
 controlnet_conditioning_scale = 0.5
 
@@ -40,6 +38,7 @@ control_modes = {
 }
 
 def load_and_convert_image(image):
+    """Load and convert images to a format that PIL can handle."""
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, bytes):
@@ -50,6 +49,7 @@ def load_and_convert_image(image):
     return image
 
 def preprocess_image(image, target_width, target_height, crop=True):
+    """Preprocess image to match the target dimensions."""
     image = load_and_convert_image(image)
     if crop:
         original_width, original_height = image.size
@@ -71,12 +71,14 @@ def preprocess_image(image, target_width, target_height, crop=True):
     return image
 
 def clear_cuda_memory():
+    """Clear CUDA memory."""
     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.ipc_collect()
 
 @spaces.GPU(duration=120)
 def generate_image(prompt, control_image, control_mode, num_steps=50, guidance=4, width=512, height=512, seed=42, random_seed=False):
+    """Generate image using the FLUX.1 ControlNet model."""
     clear_cuda_memory()
     
     if random_seed: