Update app.py

app.py

@@ -12,7 +12,6 @@ from gradio_imageslider import ImageSlider
 from PIL import Image
 from huggingface_hub import snapshot_download
 
-# Define custom CSS styling for Gradio blocks
 css = """
 #col-container {
     margin: 0 auto;
@@ -20,69 +19,52 @@ css = """
 }
 """
 
-# Determine whether GPU is available, and set the device accordingly
 if torch.cuda.is_available():
     power_device = "GPU"
     device = "cuda"
-    print("GPU is available. Using CUDA.")
 else:
     power_device = "CPU"
     device = "cpu"
-    print("GPU is not available. Using CPU.")
 
-# Get Hugging Face token from environment variables
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-print(f"Hugging Face token retrieved: {huggingface_token is not None}")
 
-# Download the model from the Hugging Face Hub
-print("Downloading model from Hugging Face Hub...")
+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,
+    token=huggingface_token, # type a new token-id.
 )
-print(f"Model downloaded to: {model_path}")
 
-# Load ControlNet model
-print("Loading ControlNet model...")
+
+# Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
 ).to(device)
-print("ControlNet model loaded.")
-
-# Load the pipeline using the downloaded model and ControlNet
-print("Loading FluxControlNetPipeline...")
 pipe = FluxControlNetPipeline.from_pretrained(
     model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
 )
 pipe.to(device)
-print("Pipeline loaded.")
 
-# Define constants for seed generation and maximum pixel budget
 MAX_SEED = 1000000
 MAX_PIXEL_BUDGET = 1024 * 1024
 
-# Function to process input image before upscaling
+
 def process_input(input_image, upscale_factor, **kwargs):
-    print(f"Processing input image with upscale factor: {upscale_factor}")
     w, h = input_image.size
     w_original, h_original = w, h
     aspect_ratio = w / h
 
     was_resized = False
 
-    # Resize the input image if the output image would exceed the pixel budget
     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."
         )
-        print("Input 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."
         )
-        # Resize the input image to fit within the maximum pixel budget
         input_image = input_image.resize(
             (
                 int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
@@ -90,18 +72,16 @@ def process_input(input_image, upscale_factor, **kwargs):
             )
         )
         was_resized = True
-        print(f"Image resized to: {input_image.size}")
 
-    # Ensure that the dimensions are multiples of 8 (required by the model)
+    # resize to multiple of 8
     w, h = input_image.size
     w = w - w % 8
     h = h - h % 8
-    print(f"Resizing image to be multiple of 8: ({w}, {h})")
 
     return input_image.resize((w, h)), w_original, h_original, was_resized
 
-# Define inference function with GPU duration hint
-@spaces.GPU(duration=42)
+
+@spaces.GPU#(duration=42)
 def infer(
     seed,
     randomize_seed,
@@ -111,109 +91,96 @@ def infer(
     controlnet_conditioning_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
-    print(f"Starting inference with seed: {seed}, randomize_seed: {randomize_seed}")
-    # Randomize the seed if the option is selected
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        print(f"Randomized seed: {seed}")
     true_input_image = input_image
-    # Process the input image for upscaling
     input_image, w_original, h_original, was_resized = process_input(
         input_image, upscale_factor
     )
-    print(f"Processed input image. Original size: ({w_original}, {h_original}), Processed size: {input_image.size}")
 
-    # Rescale the input image by the upscale factor
+    # rescale with upscale factor
     w, h = input_image.size
     control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-    print(f"Control image resized to: {control_image.size}")
 
-    # Create a random number generator with the provided seed
     generator = torch.Generator().manual_seed(seed)
 
     gr.Info("Upscaling image...")
-    print("Running the pipeline to generate output image...")
-    # Run the pipeline to generate the output image
     image = pipe(
-        prompt="",  # No specific prompt is used here
+        prompt="",
         control_image=control_image,
         controlnet_conditioning_scale=controlnet_conditioning_scale,
         num_inference_steps=num_inference_steps,
-        guidance_scale=3.5,  # Guidance scale for image generation
+        guidance_scale=3.5,
         height=control_image.size[1],
         width=control_image.size[0],
         generator=generator,
     ).images[0]
-    print("Image generation completed.")
 
-    # If the image was resized during processing, resize it back to the original target size
     if was_resized:
         gr.Info(
             f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
         )
-        print(f"Resizing output image to original target size: ({w_original * upscale_factor}, {h_original * upscale_factor})")
 
-    # Resize the generated image to the desired output size
+    # resize to target desired size
     image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
-    print(f"Final output image size: {image.size}")
     image.save("output.jpg")
-    print("Output image saved as 'output.jpg'")
-    # Return the original input image, generated image, and seed value
+    # convert to numpy
     return [true_input_image, image, seed]
 
-# Create the Gradio interface
+
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
+    # with gr.Column(elem_id="col-container"):
     gr.HTML("<center><h1>FLUX.1-Dev Upscaler</h1></center>")
     
-    # Define the button to start the upscaling process
     with gr.Row():
         run_button = gr.Button(value="Run")
 
-    # Define the input elements for the upscaling parameters
     with gr.Row():
         with gr.Column(scale=4):
-            input_im = gr.Image(label="Input Image", type="pil")  # Input image
+            input_im = gr.Image(label="Input Image", type="pil")
         with gr.Column(scale=1):
             num_inference_steps = gr.Slider(
-                label="Number of Inference Steps",  # Slider to set the number of inference steps
+                label="Number of Inference Steps",
                 minimum=8,
                 maximum=50,
                 step=1,
                 value=28,
             )
             upscale_factor = gr.Slider(
-                label="Upscale Factor",  # Slider to set the upscale factor
+                label="Upscale Factor",
                 minimum=1,
                 maximum=4,
                 step=1,
                 value=4,
             )
             controlnet_conditioning_scale = gr.Slider(
-                label="Controlnet Conditioning Scale",  # Slider for controlnet conditioning scale
+                label="Controlnet Conditioning Scale",
                 minimum=0.1,
                 maximum=1.5,
                 step=0.1,
                 value=0.6,
             )
             seed = gr.Slider(
-                label="Seed",  # Slider to set the random seed
+                label="Seed",
                 minimum=0,
                 maximum=MAX_SEED,
                 step=1,
                 value=42,
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)  # Checkbox to randomize the seed
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-    # Define the output element to display the input and output images
     with gr.Row():
         result = ImageSlider(label="Input / Output", type="pil", interactive=True)
 
-    # Define examples for users to try out
     examples = gr.Examples(
         examples=[
+        #    [42, False, "examples/image_1.jpg", 28, 4, 0.6],
             [42, False, "examples/image_2.jpg", 28, 4, 0.6],
+        #    [42, False, "examples/image_3.jpg", 28, 4, 0.6],
             [42, False, "examples/image_4.jpg", 28, 4, 0.6],
+        #    [42, False, "examples/image_5.jpg", 28, 4, 0.6],
+        #    [42, False, "examples/image_6.jpg", 28, 4, 0.6],
         ],
         inputs=[
             seed,
@@ -223,12 +190,31 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
             upscale_factor,
             controlnet_conditioning_scale,
         ],
-        fn=infer,  # Function to call for the examples
+        fn=infer,
         outputs=result,
         cache_examples="lazy",
     )
 
-    # Define the action for the run button
+    # examples = gr.Examples(
+    #     examples=[
+    #         #[42, False, "examples/image_1.jpg", 28, 4, 0.6],
+    #         [42, False, "examples/image_2.jpg", 28, 4, 0.6],
+    #         #[42, False, "examples/image_3.jpg", 28, 4, 0.6],
+    #         #[42, False, "examples/image_4.jpg", 28, 4, 0.6],
+    #         [42, False, "examples/image_5.jpg", 28, 4, 0.6],
+    #         [42, False, "examples/image_6.jpg", 28, 4, 0.6],
+    #         [42, False, "examples/image_7.jpg", 28, 4, 0.6],
+    #     ],
+    #     inputs=[
+    #         seed,
+    #         randomize_seed,
+    #         input_im,
+    #         num_inference_steps,
+    #         upscale_factor,
+    #         controlnet_conditioning_scale,
+    #     ],
+    # )
+
     gr.on(
         [run_button.click],
         fn=infer,
@@ -242,9 +228,7 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
         ],
         outputs=result,
         show_api=False,
+        # show_progress="minimal",
     )
 
-# Launch the Gradio app
-# The queue is used to handle multiple requests, sharing is disabled for privacy
-print("Launching Gradio app...")
 demo.queue().launch(share=False, show_api=False)