Update app.py

app.py

@@ -4,18 +4,13 @@ import warnings
 import os
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 from diffusers import FluxControlNetModel
 from diffusers.pipelines import FluxControlNetPipeline
 from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
-import gc
-
-# Clear memory
-gc.collect()
-if torch.cuda.is_available():
-    torch.cuda.empty_cache()
 
 css = """
 #col-container {
@@ -24,68 +19,69 @@ css = """
 }
 """
 
-# Device configuration
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float32
+if torch.cuda.is_available():
+    power_device = "GPU"
+    device = "cuda"
+else:
+    power_device = "CPU"
+    device = "cpu"
 
-huggingface_token = os.getenv("HF_TOKEN")
 
-# Modified model configuration
-model_config = {
-    "low_cpu_mem_usage": True,
-    "torch_dtype": dtype,
-    "use_safetensors": False,  # Disabled safetensors
-}
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
 model_path = snapshot_download(
-    repo_id="black-forest-labs/FLUX.1-dev",
-    repo_type="model",
-    ignore_patterns=["*.md", "*..gitattributes", "*.bin"],
-    local_dir="FLUX.1-dev", 
-    token=huggingface_token,
+    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.
 )
 
-# Load models with modified configuration
-try:
-    controlnet = FluxControlNetModel.from_pretrained(
-        "jasperai/Flux.1-dev-Controlnet-Upscaler",
-        **model_config
-    )
-    controlnet.to(device)
 
-    pipe = FluxControlNetPipeline.from_pretrained(
-        model_path,
-        controlnet=controlnet,
-        **model_config
-    )
-    pipe.to(device)
+# Load pipeline
+controlnet = FluxControlNetModel.from_pretrained(
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
+).to(device)
+pipe = FluxControlNetPipeline.from_pretrained(
+    model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
+)
+pipe.to(device)
 
-except Exception as e:
-    print(f"Error loading models: {str(e)}")
-    raise
+MAX_SEED = 1000000
+MAX_PIXEL_BUDGET = 1024 * 1024
 
-# Enable optimizations
-pipe.enable_attention_slicing(1)
-pipe.enable_vae_slicing()
 
-MAX_SEED = 1000000
-MAX_PIXEL_BUDGET = 64 * 64
+def process_input(input_image, upscale_factor, **kwargs):
+    w, h = input_image.size
+    w_original, h_original = w, h
+    aspect_ratio = w / h
+
+    was_resized = False
+
+    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."
+        )
+        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."
+        )
+        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),
+            )
+        )
+        was_resized = True
 
-def process_input(input_image, upscale_factor):
-    input_image = input_image.convert('RGB')
-    
+    # resize to multiple of 8
     w, h = input_image.size
-    max_size = int(np.sqrt(MAX_PIXEL_BUDGET))
-    
-    new_w = min(w, max_size)
-    new_h = min(h, max_size)
-    input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
-    
-    w = new_w - new_w % 8
-    h = new_h - new_h % 8
-    
-    return input_image.resize((w, h)), w, h
+    w = w - w % 8
+    h = h - h % 8
+
+    return input_image.resize((w, h)), w_original, h_original, was_resized
 
+
+@spaces.GPU#(duration=42)
 def infer(
     seed,
     randomize_seed,
@@ -95,38 +91,42 @@ def infer(
     controlnet_conditioning_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
-    try:
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-            
-        input_image, w, h = process_input(input_image, upscale_factor)
-        
-        with torch.inference_mode():
-            generator = torch.Generator(device=device).manual_seed(seed)
-            image = pipe(
-                prompt="",
-                control_image=input_image,
-                controlnet_conditioning_scale=controlnet_conditioning_scale,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=1.5,
-                height=h,
-                width=w,
-                generator=generator,
-            ).images[0]
-            
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        return [input_image, image, seed]
-
-    except Exception as e:
-        gr.Error(f"Error: {str(e)}")
-        return None
+    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."
+        )
+
+    # 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]
+
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     with gr.Row():