Update app.py

app.py

@@ -19,37 +19,53 @@ css = """
 }
 """
 
+# Device and dtype setup
 if torch.cuda.is_available():
     power_device = "GPU"
-    device = "cuda"
+    device = "cuda" 
+    dtype = torch.bfloat16
 else:
     power_device = "CPU"
     device = "cpu"
-
+    dtype = torch.float32
 
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
 model_path = snapshot_download(
-    repo_id="black-forest-labs/FLUX.1-dev", 
-    repo_type="model", 
+    repo_id="black-forest-labs/FLUX.1-dev",
+    repo_type="model",
     ignore_patterns=["*.md", "*..gitattributes"],
-    local_dir="FLUX.1-dev",
-    token=huggingface_token, # type a new token-id.
+    local_dir="FLUX.1-dev", 
+    token=huggingface_token,
 )
 
-
-# Load pipeline
+# Load pipeline with memory optimizations
 controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    torch_dtype=dtype
 ).to(device)
+
 pipe = FluxControlNetPipeline.from_pretrained(
-    model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
+    model_path,
+    controlnet=controlnet,
+    torch_dtype=dtype
 )
 pipe.to(device)
 
+# Enable memory optimizations
+pipe.enable_model_cpu_offload()
+pipe.enable_attention_slicing()
+
 MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 1024 * 1024
+MAX_PIXEL_BUDGET = 512 * 512  # Reduced from 1024 * 1024
 
+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
+            return False
+    return True
 
 def process_input(input_image, upscale_factor, **kwargs):
     w, h = input_image.size
@@ -80,8 +96,7 @@ def process_input(input_image, upscale_factor, **kwargs):
 
     return input_image.resize((w, h)), w_original, h_original, was_resized
 
-
-@spaces.GPU#(duration=42)
+@spaces.GPU
 def infer(
     seed,
     randomize_seed,
@@ -91,46 +106,60 @@ def infer(
     controlnet_conditioning_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    true_input_image = input_image
-    input_image, w_original, h_original, was_resized = process_input(
-        input_image, upscale_factor
-    )
-
-    # rescale with upscale factor
-    w, h = input_image.size
-    control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-    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]
-
-    if was_resized:
-        gr.Info(
-            f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
+    try:
+        if not check_resources():
+            gr.Warning("System resources are running low. Try reducing parameters.")
+            return None
+
+        if device == "cuda":
+            torch.cuda.empty_cache()
+
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+            
+        true_input_image = input_image
+        input_image, w_original, h_original, was_resized = process_input(
+            input_image, upscale_factor
         )
 
-    # resize to target desired size
-    image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
-    image.save("output.jpg")
-    # convert to numpy
-    return [true_input_image, image, seed]
+        # rescale with upscale factor
+        w, h = input_image.size
+        control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
+
+        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]
+
+        if was_resized:
+            gr.Info(
+                f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} 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]
 
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-
+    except RuntimeError as e:
+        if "out of memory" in str(e):
+            gr.Warning("Not enough GPU memory. Try reducing the upscale factor or image size.")
+            return None
+        raise e
+    except Exception as e:
+        gr.Error(f"An error occurred: {str(e)}")
+        return None
 
+with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     with gr.Row():
         run_button = gr.Button(value="Run")
 
@@ -148,9 +177,9 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
             upscale_factor = gr.Slider(
                 label="Upscale Factor",
                 minimum=1,
-                maximum=4,
+                maximum=2,  # Reduced from 4
                 step=1,
-                value=4,
+                value=2,    # Reduced default
             )
             controlnet_conditioning_scale = gr.Slider(
                 label="Controlnet Conditioning Scale",
@@ -174,9 +203,8 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
 
     examples = gr.Examples(
         examples=[
-            [42, False, "z1.webp", 28, 4, 0.6],
-            [42, False, "z2.webp", 28, 4, 0.6],
-
+            [42, False, "z1.webp", 28, 2, 0.6],  # Updated upscale factor
+            [42, False, "z2.webp", 28, 2, 0.6],  # Updated upscale factor
         ],
         inputs=[
             seed,
@@ -204,7 +232,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         ],
         outputs=result,
         show_api=False,
-        # show_progress="minimal",
     )
 
-demo.queue().launch(share=False)
+demo.queue().launch(share=False)