Update app.py

app.py

@@ -1,776 +1,3 @@
-# import logging
-# import random
-# import warnings
-# import gradio as gr
-# import os
-# import shutil
-# import subprocess
-# import spaces
-# import torch
-# import numpy as np
-# from diffusers import FluxControlNetModel
-# from diffusers.pipelines import FluxControlNetPipeline
-# from PIL import Image
-# from huggingface_hub import snapshot_download, login
-# import io
-# import base64
-# from flask import Flask, request, jsonify
-# from concurrent.futures import ThreadPoolExecutor
-# from flask_cors import CORS
-# import threading
-
-# # Configure logging
-# logging.basicConfig(level=logging.INFO)
-# logger = logging.getLogger(__name__)
-
-# app = Flask(__name__)
-# CORS(app)
-
-# # Function to check disk usage
-# def check_disk_space():
-#     result = subprocess.run(['df', '-h'], capture_output=True, text=True)
-#     logger.info("Disk space usage:\n%s", result.stdout)
-
-# # Function to clear Hugging Face cache
-# def clear_huggingface_cache():
-#     cache_dir = os.path.expanduser('~/.cache/huggingface')
-#     if os.path.exists(cache_dir):
-#         shutil.rmtree(cache_dir)  # Removes the entire cache directory
-#         logger.info("Cleared Hugging Face cache at: %s", cache_dir)
-#     else:
-#         logger.info("No Hugging Face cache found.")
-
-# # Check disk space
-# check_disk_space()
-
-# # Clear Hugging Face cache
-# clear_huggingface_cache()
-
-# # Add config to store base64 images
-# app.config['image_outputs'] = {}
-
-# # ThreadPoolExecutor for managing image processing threads
-# executor = ThreadPoolExecutor()
-
-# # Determine the device (GPU or CPU)
-# if torch.cuda.is_available():
-#     device = "cuda"
-#     logger.info("CUDA is available. Using GPU.")
-# else:
-#     device = "cpu"
-#     logger.info("CUDA is not available. Using CPU.")
-
-# # Load model from Huggingface Hub
-# huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-# if huggingface_token:
-#     login(token=huggingface_token)
-#     logger.info("Hugging Face token found and logged in.")
-# else:
-#     logger.warning("Hugging Face token not found in environment variables.")
-
-# logger.info("Hugging Face token: %s", huggingface_token)
-
-# # Download model using snapshot_download
-
-# model_path = snapshot_download(
-#         repo_id="black-forest-labs/FLUX.1-dev", 
-#         repo_type="model", 
-#         ignore_patterns=["*.md", "*..gitattributes"],
-#         local_dir="FLUX.1-dev",
-#         token=huggingface_token)
-# logger.info("Model downloaded to: %s", model_path)
-
-# # Load pipeline
-# logger.info('Loading ControlNet model.')
-
-# controlnet = FluxControlNetModel.from_pretrained(
-#         "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-#     ).to(device)
-# logger.info("ControlNet model loaded successfully.")
-
-# logger.info('Loading pipeline.')
-
-# pipe = FluxControlNetPipeline.from_pretrained(
-#         model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
-#     ).to(device)
-# logger.info("Pipeline loaded successfully.")
-
-# MAX_SEED = 1000000
-# MAX_PIXEL_BUDGET = 1024 * 1024
-
-
-# @spaces.GPU
-# def process_input(input_image, upscale_factor):
-#     w, h = input_image.size
-#     aspect_ratio = w / h
-#     was_resized = False
-
-#     # Resize if input size exceeds the maximum pixel budget
-#     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-#         warnings.warn(f"Requested output image is too large. Resizing to fit within pixel 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
-
-#     # Adjust dimensions to be a multiple of 8
-#     w, h = input_image.size
-#     w = w - w % 8
-#     h = h - h % 8
-
-#     return input_image.resize((w, h)), was_resized
-
-# @spaces.GPU
-# def run_inference(process_id, input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale):
-#     logger.info("Processing inference for process_id: %s", process_id)
-#     input_image, was_resized = process_input(input_image, upscale_factor)
-
-#     # Rescale image for ControlNet processing
-#     w, h = input_image.size
-#     control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-#     # Set the random generator for inference
-#     generator = torch.Generator().manual_seed(seed)
-
-#     # Perform inference using the pipeline
-#     logger.info("Running pipeline for process_id: %s", process_id)
-#     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]
-
-#     # Resize output image back to the original dimensions if needed
-#     if was_resized:
-#         original_size = (input_image.width * upscale_factor, input_image.height * upscale_factor)
-#         image = image.resize(original_size)
-
-#     # Convert the output image to base64
-#     buffered = io.BytesIO()
-#     image.save(buffered, format="JPEG")
-#     image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-#     # Store the result in the shared dictionary
-#     app.config['image_outputs'][process_id] = image_base64
-#     logger.info("Inference completed for process_id: %s", process_id)
-
-# @app.route('/infer', methods=['POST'])
-# def infer():
-#     # Check if the file was provided in the form-data
-#     if 'input_image' not in request.files:
-#         logger.error("No image file provided in request.")
-#         return jsonify({
-#             "status": "error",
-#             "message": "No input_image file provided"
-#         }), 400
-
-#     # Get the uploaded image file from the request
-#     file = request.files['input_image']
-
-#     # Check if a file was uploaded
-#     if file.filename == '':
-#         logger.error("No selected file in form-data.")
-#         return jsonify({
-#             "status": "error",
-#             "message": "No selected file"
-#         }), 400
-
-#     # Convert the image to Base64 for internal processing
-#     input_image = Image.open(file)
-#     buffered = io.BytesIO()
-#     input_image.save(buffered, format="JPEG")
-
-#     # Retrieve additional parameters from the request (if any)
-#     seed = request.form.get("seed", 42, type=int)
-#     randomize_seed = request.form.get("randomize_seed", 'true').lower() == 'true'
-#     num_inference_steps = request.form.get("num_inference_steps", 28, type=int)
-#     upscale_factor = request.form.get("upscale_factor", 4, type=int)
-#     controlnet_conditioning_scale = request.form.get("controlnet_conditioning_scale", 0.6, type=float)
-
-#     # Randomize seed if specified
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-#         logger.info("Seed randomized to: %d", seed)
-
-#     # Create a unique process ID for this request
-#     process_id = str(random.randint(1000, 9999))
-#     logger.info("Process started with process_id: %s", process_id)
-
-#     # Set the status to 'in_progress'
-#     app.config['image_outputs'][process_id] = None
-
-#     # Run the inference in a separate thread
-#     executor.submit(run_inference, process_id, input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale)
-
-#     # Return the process ID
-#     return jsonify({
-#         "process_id": process_id,
-#         "message": "Processing started"
-#     })
-
-
-# # Modify status endpoint to receive process_id in request body
-# @app.route('/status', methods=['GET'])
-# def status():
-#     # Get the process_id from the query parameters
-#     process_id = request.args.get('process_id')
-
-#     # Check if process_id was provided
-#     if not process_id:
-#         logger.error("Process ID not provided in request.")
-#         return jsonify({
-#             "status": "error",
-#             "message": "Process ID is required"
-#         }), 400
-
-#     # Check if the process_id exists in the dictionary
-#     if process_id not in app.config['image_outputs']:
-#         logger.error("Invalid process ID: %s", process_id)
-#         return jsonify({
-#             "status": "error",
-#             "message": "Invalid process ID"
-#         }), 404
-
-#     # Check the status of the image processing
-#     image_base64 = app.config['image_outputs'][process_id]
-#     if image_base64 is None:
-#         logger.info("Process ID %s is still in progress.", process_id)
-#         return jsonify({
-#             "status": "in_progress"
-#         })
-#     else:
-#         logger.info("Process ID %s completed successfully.", process_id)
-#         return jsonify({
-#             "status": "completed",
-#             "output_image": image_base64
-#         })
-
-
-# if __name__ == '__main__':
-#     app.run(debug=True,host='0.0.0.0')
-
-
-
-
-
-# import logging
-# import random
-# import warnings
-# import gradio as gr
-# import os
-# import shutil
-# import subprocess
-# import torch
-# import numpy as np
-# from diffusers import FluxControlNetModel
-# from diffusers.pipelines import FluxControlNetPipeline
-# from PIL import Image
-# from huggingface_hub import snapshot_download, login
-# import io
-# import base64
-# import threading
-
-# # Configure logging
-# logging.basicConfig(level=logging.INFO)
-# logger = logging.getLogger(__name__)
-
-# # Function to check disk usage
-# def check_disk_space():
-#     result = subprocess.run(['df', '-h'], capture_output=True, text=True)
-#     logger.info("Disk space usage:\n%s", result.stdout)
-
-# # Function to clear Hugging Face cache
-# def clear_huggingface_cache():
-#     cache_dir = os.path.expanduser('~/.cache/huggingface')
-#     if os.path.exists(cache_dir):
-#         shutil.rmtree(cache_dir)  # Removes the entire cache directory
-#         logger.info("Cleared Hugging Face cache at: %s", cache_dir)
-#     else:
-#         logger.info("No Hugging Face cache found.")
-
-# # Check disk space
-# check_disk_space()
-
-# # Clear Hugging Face cache
-# clear_huggingface_cache()
-
-# # Add config to store base64 images
-# image_outputs = {}
-
-# # Determine the device (GPU or CPU)
-# if torch.cuda.is_available():
-#     device = "cuda"
-#     logger.info("CUDA is available. Using GPU.")
-# else:
-#     device = "cpu"
-#     logger.info("CUDA is not available. Using CPU.")
-
-# # Load model from Huggingface Hub
-# huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-# if huggingface_token:
-#     login(token=huggingface_token)
-#     logger.info("Hugging Face token found and logged in.")
-# else:
-#     logger.warning("Hugging Face token not found in environment variables.")
-
-# # Download model using snapshot_download
-# model_path = snapshot_download(
-#     repo_id="black-forest-labs/FLUX.1-dev", 
-#     repo_type="model", 
-#     ignore_patterns=["*.md", "*..gitattributes"],
-#     local_dir="FLUX.1-dev",
-#     token=huggingface_token
-# )
-# logger.info("Model downloaded to: %s", model_path)
-
-# # Load pipeline
-# logger.info('Loading ControlNet model.')
-# controlnet = FluxControlNetModel.from_pretrained(
-#     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("ControlNet model loaded successfully.")
-
-# logger.info('Loading pipeline.')
-# pipe = FluxControlNetPipeline.from_pretrained(
-#     model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("Pipeline loaded successfully.")
-
-# MAX_SEED = 1000000
-# MAX_PIXEL_BUDGET = 1024 * 1024
-
-# def process_input(input_image, upscale_factor):
-#     w, h = input_image.size
-#     aspect_ratio = w / h
-#     was_resized = False
-
-#     # Resize if input size exceeds the maximum pixel budget
-#     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-#         warnings.warn(f"Requested output image is too large. Resizing to fit within pixel 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
-
-#     # Adjust dimensions to be a multiple of 8
-#     w, h = input_image.size
-#     w = w - w % 8
-#     h = h - h % 8
-
-#     return input_image.resize((w, h)), was_resized
-
-# def run_inference(input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale):
-#     logger.info("Running inference")
-#     input_image, was_resized = process_input(input_image, upscale_factor)
-
-#     # Rescale image for ControlNet processing
-#     w, h = input_image.size
-#     control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-#     # Set the random generator for inference
-#     generator = torch.Generator().manual_seed(seed)
-
-#     # Perform inference using the pipeline
-#     logger.info("Running pipeline")
-#     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]
-
-#     # Resize output image back to the original dimensions if needed
-#     if was_resized:
-#         original_size = (input_image.width * upscale_factor, input_image.height * upscale_factor)
-#         image = image.resize(original_size)
-
-#     # Convert the output image to base64
-#     buffered = io.BytesIO()
-#     image.save(buffered, format="JPEG")
-#     image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-#     logger.info("Inference completed")
-#     return image_base64
-
-# # Define Gradio interface
-# def gradio_interface(input_image, upscale_factor=4, seed=42, num_inference_steps=28, controlnet_conditioning_scale=0.6):
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-#         logger.info("Seed randomized to: %d", seed)
-
-#     # Run inference
-#     output_image_base64 = run_inference(input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale)
-    
-#     return Image.open(io.BytesIO(base64.b64decode(output_image_base64)))
-
-# # Create Gradio interface
-# iface = gr.Interface(
-#     fn=gradio_interface,
-#     inputs=[
-#         gr.Image(type="pil", label="Input Image"),
-#         gr.Slider(min=1, max=8, step=1, label="Upscale Factor"),
-#         gr.Slider(min=0, max=MAX_SEED, step=1, label="Seed"),
-#         gr.Slider(min=1, max=100, step=1, label="Inference Steps"),
-#         gr.Slider(min=0.0, max=1.0, step=0.1, label="ControlNet Conditioning Scale")
-#     ],
-#     outputs=gr.Image(label="Output Image"),
-#     title="ControlNet Image Upscaling",
-#     description="Upload an image to upscale using the ControlNet model."
-# )
-
-# # Launch Gradio app
-# if __name__ == '__main__':
-#     iface.launch()
-
-
-# import logging
-# import random
-# import warnings
-# import gradio as gr
-# import os
-# import shutil
-# import spaces
-# import subprocess
-# import torch
-# import numpy as np
-# from diffusers import FluxControlNetModel
-# from diffusers.pipelines import FluxControlNetPipeline
-# from PIL import Image
-# from huggingface_hub import snapshot_download, login
-# import io
-# import base64
-# from concurrent.futures import ThreadPoolExecutor
-
-# # Configure logging
-# logging.basicConfig(level=logging.INFO)
-# logger = logging.getLogger(__name__)
-
-# # ThreadPoolExecutor for managing image processing threads
-# executor = ThreadPoolExecutor()
-
-# # Determine the device (GPU or CPU)
-# if torch.cuda.is_available():
-#     device = "cuda"
-#     logger.info("CUDA is available. Using GPU.")
-# else:
-#     device = "cpu"
-#     logger.info("CUDA is not available. Using CPU.")
-
-# # Load model from Huggingface Hub
-# huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-# if huggingface_token:
-#     login(token=huggingface_token)
-#     logger.info("Hugging Face token found and logged in.")
-# else:
-#     logger.warning("Hugging Face token not found in environment variables.")
-
-# # Download model using snapshot_download
-# model_path = snapshot_download(
-#     repo_id="black-forest-labs/FLUX.1-dev",
-#     repo_type="model",
-#     ignore_patterns=["*.md", "*..gitattributes"],
-#     local_dir="FLUX.1-dev",
-#     token=huggingface_token
-# )
-# logger.info("Model downloaded to: %s", model_path)
-
-# # Load pipeline
-# logger.info('Loading ControlNet model.')
-# controlnet = FluxControlNetModel.from_pretrained(
-#     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("ControlNet model loaded successfully.")
-
-# logger.info('Loading pipeline.')
-# pipe = FluxControlNetPipeline.from_pretrained(
-#     model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("Pipeline loaded successfully.")
-
-# MAX_SEED = 1000000
-# MAX_PIXEL_BUDGET = 1024 * 1024
-
-# @spaces.GPU
-# def process_input(input_image, upscale_factor):
-#     w, h = input_image.size
-#     aspect_ratio = w / h
-#     was_resized = False
-
-#     # Resize if input size exceeds the maximum pixel budget
-#     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-#         warnings.warn(f"Requested output image is too large. Resizing to fit within pixel 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
-
-#     # Adjust dimensions to be a multiple of 8
-#     w, h = input_image.size
-#     w = w - w % 8
-#     h = h - h % 8
-
-#     return input_image.resize((w, h)), was_resized
-    
-
-# @spaces.GPU
-# def run_inference(input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale):
-#     logger.info("Processing inference.")
-#     input_image, was_resized = process_input(input_image, upscale_factor)
-
-#     # Rescale image for ControlNet processing
-#     w, h = input_image.size
-#     control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-#     # Set the random generator for inference
-#     generator = torch.Generator().manual_seed(seed)
-
-#     # Perform inference using the pipeline
-#     logger.info("Running pipeline.")
-#     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]
-
-#     # Resize output image back to the original dimensions if needed
-#     if was_resized:
-#         original_size = (input_image.width * upscale_factor, input_image.height * upscale_factor)
-#         image = image.resize(original_size)
-
-#     return image
-
-# def run_gradio_app():
-#     with gr.Blocks() as app:
-#         gr.Markdown("## Image Upscaler using ControlNet")
-
-#         # Define the inputs and outputs
-#         input_image = gr.Image(type="pil", label="Input Image")
-#         upscale_factor = gr.Slider(minimum=1, maximum=8, step=1, label="Upscale Factor")
-#         seed = gr.Slider(minimum=0, maximum=100, step=1, label="Seed")
-#         num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Inference Steps")
-#         controlnet_conditioning_scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="ControlNet Conditioning Scale")
-
-#         output_image = gr.Image(type="pil", label="Output Image")
-
-#         # Create a button to trigger the processing
-#         submit_button = gr.Button("Upscale Image")
-
-#         # Define the function to run when the button is clicked
-#         submit_button.click(run_inference, 
-#                              inputs=[input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale], 
-#                              outputs=output_image)
-
-#     app.launch()
-
-# if __name__ == "__main__":
-#     run_gradio_app()
-
-
-
-
-# import logging
-# import random
-# import warnings
-# import gradio as gr
-# import os
-# import shutil,spaces
-# import subprocess
-# import torch
-# import numpy as np
-# from diffusers import FluxControlNetModel
-# from diffusers.pipelines import FluxControlNetPipeline
-# from PIL import Image
-# from huggingface_hub import snapshot_download, login
-# import io
-# import base64
-# from fastapi import FastAPI, File, UploadFile
-# from fastapi.responses import JSONResponse
-# from fastapi.middleware.cors import CORSMiddleware
-# from concurrent.futures import ThreadPoolExecutor
-
-# # Configure logging
-# logging.basicConfig(level=logging.INFO)
-# logger = logging.getLogger(__name__)
-
-# # FastAPI app for image processing
-# app = FastAPI()
-# app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
-
-# # ThreadPoolExecutor for managing image processing threads
-# executor = ThreadPoolExecutor()
-
-# # Determine the device (GPU or CPU)
-# if torch.cuda.is_available():
-#     device = "cuda"
-#     logger.info("CUDA is available. Using GPU.")
-# else:
-#     device = "cpu"
-#     logger.info("CUDA is not available. Using CPU.")
-
-# # Load model from Huggingface Hub
-# huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-# if huggingface_token:
-#     login(token=huggingface_token)
-#     logger.info("Hugging Face token found and logged in.")
-# else:
-#     logger.warning("Hugging Face token not found in environment variables.")
-
-# # Download model using snapshot_download
-# model_path = snapshot_download(
-#     repo_id="black-forest-labs/FLUX.1-dev",
-#     repo_type="model",
-#     ignore_patterns=["*.md", "*..gitattributes"],
-#     local_dir="FLUX.1-dev",
-#     token=huggingface_token
-# )
-# logger.info("Model downloaded to: %s", model_path)
-
-# # Load pipeline
-# logger.info('Loading ControlNet model.')
-# controlnet = FluxControlNetModel.from_pretrained(
-#     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("ControlNet model loaded successfully.")
-
-# logger.info('Loading pipeline.')
-# pipe = FluxControlNetPipeline.from_pretrained(
-#     model_path, controlnet=controlnet, torch_dtype=torch.bfloat16
-# ).to(device)
-# logger.info("Pipeline loaded successfully.")
-
-# MAX_SEED = 1000000
-# MAX_PIXEL_BUDGET = 1024 * 1024
-
-# @spaces.GPU
-# def process_input(input_image, upscale_factor):
-#     w, h = input_image.size
-#     aspect_ratio = w / h
-#     was_resized = False
-
-#     # Resize if input size exceeds the maximum pixel budget
-#     if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
-#         warnings.warn(f"Requested output image is too large. Resizing to fit within pixel 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
-
-#     # Adjust dimensions to be a multiple of 8
-#     w, h = input_image.size
-#     w = w - w % 8
-#     h = h - h % 8
-
-#     return input_image.resize((w, h)), was_resized
-
-# @spaces.GPU
-# def run_inference(input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale):
-#     logger.info("Processing inference.")
-#     input_image, was_resized = process_input(input_image, upscale_factor)
-
-#     # Rescale image for ControlNet processing
-#     w, h = input_image.size
-#     control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-#     # Set the random generator for inference
-#     generator = torch.Generator().manual_seed(seed)
-
-#     # Perform inference using the pipeline
-#     logger.info("Running pipeline.")
-#     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]
-
-#     # Resize output image back to the original dimensions if needed
-#     if was_resized:
-#         original_size = (input_image.width * upscale_factor, input_image.height * upscale_factor)
-#         image = image.resize(original_size)
-
-#     # Convert the output image to base64
-#     buffered = io.BytesIO()
-#     image.save(buffered, format="JPEG")
-#     image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-#     return image_base64
-
-# @app.post("/infer")
-# async def infer(input_image: UploadFile = File(...), 
-#                  upscale_factor: int = 4, 
-#                  seed: int = 42, 
-#                  num_inference_steps: int = 28, 
-#                  controlnet_conditioning_scale: float = 0.6):
-#     logger.info("Received request for inference.")
-    
-#     # Read the uploaded image
-#     contents = await input_image.read()
-#     image = Image.open(io.BytesIO(contents))
-
-#     # Run inference in a separate thread
-#     base64_image = await executor.submit(run_inference, image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale)
-
-#     return JSONResponse(content={"base64_image": base64_image})
-
-# def run_gradio_app():
-#     with gr.Blocks() as app:
-#         gr.Markdown("## Image Upscaler using ControlNet")
-
-#         # Define the inputs and outputs
-#         input_image = gr.Image(type="pil", label="Input Image")
-#         upscale_factor = gr.Slider(minimum=1, maximum=8, step=1, label="Upscale Factor")
-#         seed = gr.Slider(minimum=0, maximum=100, step=1, label="Seed")
-#         num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Inference Steps")
-#         controlnet_conditioning_scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="ControlNet Conditioning Scale")
-
-#         output_image = gr.Image(type="pil", label="Output Image")
-#         output_base64 = gr.Textbox(label="Base64 String", interactive=False)
-
-#         # Create a button to trigger the processing
-#         submit_button = gr.Button("Upscale Image")
-
-#         # Define the function to run when the button is clicked
-#         submit_button.click(run_inference, 
-#                              inputs=[input_image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale], 
-#                              outputs=[output_image, output_base64])
-
-#     app.launch()
-
-# if __name__ == "__main__":
-#     run_gradio_app()
-
-
-
 import spaces
 import logging
 import random
@@ -792,6 +19,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from concurrent.futures import ThreadPoolExecutor
 import uvicorn
 import asyncio
+import time  # Import time module for measuring execution time
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -915,6 +143,9 @@ async def infer(input_image: UploadFile = File(...),
                  controlnet_conditioning_scale: float = 0.6):
     logger.info("Received request for inference.")
     
+    # Start timing the entire inference process
+    start_time = time.time()
+
     # Read the uploaded image
     contents = await input_image.read()
     image = Image.open(io.BytesIO(contents))
@@ -925,13 +156,11 @@ async def infer(input_image: UploadFile = File(...),
     # Run inference in a separate thread
     base64_image = await loop.run_in_executor(executor, run_inference, image, upscale_factor, seed, num_inference_steps, controlnet_conditioning_scale)
 
-    return JSONResponse(content={"base64_image": base64_image})
+    # Calculate the time taken
+    time_taken = time.time() - start_time
+
+    return JSONResponse(content={"base64_image": base64_image, "time_taken": time_taken})
 
 if __name__ == "__main__":
     # Start FastAPI server
     uvicorn.run(app, host="0.0.0.0", port=7860)
-
-
-
-
-