Update app.py

app.py

@@ -591,12 +591,193 @@
 
 
 
+# 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 logging
 import random
 import warnings
 import gradio as gr
 import os
-import shutil,spaces
+import shutil
 import subprocess
 import torch
 import numpy as np
@@ -610,6 +791,8 @@ from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from concurrent.futures import ThreadPoolExecutor
+import uvicorn
+import spaces
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -672,7 +855,7 @@ def process_input(input_image, upscale_factor):
 
     # 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.")
+        warnings.warn("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),
@@ -753,7 +936,6 @@ def run_gradio_app():
         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
@@ -762,12 +944,19 @@ def run_gradio_app():
         # 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])
+                             outputs=[output_base64])
 
     app.launch()
 
 if __name__ == "__main__":
-    run_gradio_app()
+    # Run Gradio app in a separate thread
+    import threading
+    gradio_thread = threading.Thread(target=run_gradio_app)
+    gradio_thread.start()
+
+    # Start FastAPI server
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+