Update app.py

app.py

@@ -19,17 +19,22 @@ css = """
 }
 """
 
-# Device and dtype setup
+# Device and dtype setup with lower precision
 if torch.cuda.is_available():
     power_device = "GPU"
     device = "cuda" 
-    dtype = torch.bfloat16
+    dtype = torch.float16  # Changed to float16 for less memory usage
 else:
     power_device = "CPU"
     device = "cpu"
     dtype = torch.float32
 
-huggingface_token = os.getenv("HUGGINFACE_TOKEN")
+# Reduce CUDA memory usage
+torch.cuda.empty_cache()
+if torch.cuda.is_available():
+    torch.cuda.set_per_process_memory_fraction(0.7)  # Use only 70% of GPU memory
+
+huggingface_token = os.getenv("HF_TOKEN")
 
 model_path = snapshot_download(
     repo_id="black-forest-labs/FLUX.1-dev",
@@ -39,35 +44,44 @@ model_path = snapshot_download(
     token=huggingface_token,
 )
 
-# Load pipeline with memory optimizations
+# Load pipeline with more memory optimizations
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler",
-    torch_dtype=dtype
+    torch_dtype=dtype,
+    low_cpu_mem_usage=True,
+    use_safetensors=True
 ).to(device)
 
 pipe = FluxControlNetPipeline.from_pretrained(
     model_path,
     controlnet=controlnet,
-    torch_dtype=dtype
+    torch_dtype=dtype,
+    low_cpu_mem_usage=True,
+    use_safetensors=True
 )
-pipe.to(device)
 
-# Enable memory optimizations
+# Enable all possible memory optimizations
 pipe.enable_model_cpu_offload()
-pipe.enable_attention_slicing()
+pipe.enable_attention_slicing(1)
+pipe.enable_sequential_cpu_offload()
+pipe.enable_vae_slicing()
 
+# Further reduce memory usage
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 512 * 512  # Reduced from 1024 * 1024
+MAX_PIXEL_BUDGET = 256 * 256  # Further reduced from 512 * 512
 
 def check_resources():
     if torch.cuda.is_available():
         gpu_memory = torch.cuda.get_device_properties(0).total_memory
         memory_allocated = torch.cuda.memory_allocated(0)
-        if memory_allocated/gpu_memory > 0.9:  # 90% threshold
+        if memory_allocated/gpu_memory > 0.8:  # 80% threshold
             return False
     return True
 
 def process_input(input_image, upscale_factor, **kwargs):
+    # Convert image to RGB mode to ensure compatibility
+    input_image = input_image.convert('RGB')
+    
     w, h = input_image.size
     w_original, h_original = w, h
     aspect_ratio = w / h
@@ -76,16 +90,17 @@ def process_input(input_image, upscale_factor, **kwargs):
 
     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
         warnings.warn(
-            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels."
+            f"Requested output image is too large. Resizing..."
         )
         gr.Info(
-            f"Requested output image is too large ({w * upscale_factor}x{h * upscale_factor}). Resizing input to ({int(aspect_ratio * MAX_PIXEL_BUDGET ** 0.5 // upscale_factor), int(MAX_PIXEL_BUDGET ** 0.5 // aspect_ratio // upscale_factor)}) pixels budget."
+            f"Resizing input image to fit memory constraints..."
         )
         input_image = input_image.resize(
             (
                 int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
                 int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
-            )
+            ),
+            Image.LANCZOS
         )
         was_resized = True
 
@@ -111,6 +126,7 @@ def infer(
             gr.Warning("System resources are running low. Try reducing parameters.")
             return None
 
+        # Clear CUDA cache before processing
         if device == "cuda":
             torch.cuda.empty_cache()
 
@@ -129,25 +145,25 @@ def infer(
         generator = torch.Generator().manual_seed(seed)
 
         gr.Info("Upscaling image...")
-        image = pipe(
-            prompt="",
-            control_image=control_image,
-            controlnet_conditioning_scale=controlnet_conditioning_scale,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=3.5,
-            height=control_image.size[1],
-            width=control_image.size[0],
-            generator=generator,
-        ).images[0]
+        with torch.inference_mode():  # Use inference mode to save memory
+            image = pipe(
+                prompt="",
+                control_image=control_image,
+                controlnet_conditioning_scale=controlnet_conditioning_scale,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=3.5,
+                height=control_image.size[1],
+                width=control_image.size[0],
+                generator=generator,
+            ).images[0]
 
         if was_resized:
             gr.Info(
-                f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
+                f"Resizing output image to final size..."
             )
 
         # resize to target desired size
         image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
-        image.save("output.jpg")
         return [true_input_image, image, seed]
 
     except RuntimeError as e:
@@ -170,23 +186,23 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
             num_inference_steps = gr.Slider(
                 label="Number of Inference Steps",
                 minimum=8,
-                maximum=50,
+                maximum=30,  # Reduced from 50
                 step=1,
-                value=28,
+                value=20,    # Reduced from 28
             )
             upscale_factor = gr.Slider(
                 label="Upscale Factor",
                 minimum=1,
-                maximum=2,  # Reduced from 4
+                maximum=2,
                 step=1,
-                value=2,    # Reduced default
+                value=1,    # Reduced default
             )
             controlnet_conditioning_scale = gr.Slider(
                 label="Controlnet Conditioning Scale",
                 minimum=0.1,
-                maximum=1.5,
+                maximum=1.0,  # Reduced from 1.5
                 step=0.1,
-                value=0.6,
+                value=0.5,   # Reduced from 0.6
             )
             seed = gr.Slider(
                 label="Seed",
@@ -205,8 +221,8 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     
     examples = gr.Examples(
         examples=[
-            [42, False, os.path.join(current_dir, "z1.webp"), 28, 2, 0.6],
-            [42, False, os.path.join(current_dir, "z2.webp"), 28, 2, 0.6],
+            [42, False, os.path.join(current_dir, "z1.webp"), 20, 1, 0.5],  # Reduced parameters
+            [42, False, os.path.join(current_dir, "z2.webp"), 20, 1, 0.5],  # Reduced parameters
         ],
         inputs=[
             seed,
@@ -236,4 +252,11 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         show_api=False,
     )
 
-demo.queue().launch(share=False)
+# Launch with minimal memory usage
+demo.queue(max_size=1).launch(
+    share=False,
+    debug=True,
+    show_error=True,
+    max_threads=1,
+    enable_queue=True
+)