Update app.py

app.py

@@ -7,6 +7,11 @@ import tempfile
 from PIL import Image
 from huggingface_hub import hf_hub_download, login
 import torch
+import logging
+
+# --- Logging Setup ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # --- Configuration ---
 LORA_PATH = "ckpt/SR_LoRA/model_20001.pkl"
@@ -21,26 +26,26 @@ CHECKPOINT_FILES_CONFIG = {
 
 # --- Device Detection ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {DEVICE}")
+logger.info(f"Using device: {DEVICE}")
 
 # --- Hugging Face Token ---
 HF_AUTH_TOKEN = os.environ.get("HF_TOKEN")
 if HF_AUTH_TOKEN:
     try:
         login(token=HF_AUTH_TOKEN)
-        print("Successfully logged in to Hugging Face Hub.")
+        logger.info("Successfully logged in to Hugging Face Hub.")
     except Exception as e:
-        print(f"Warning: Hugging Face login failed: {e}")
+        logger.warning(f"Hugging Face login failed: {e}")
 else:
-    print("Warning: HF_TOKEN not found. Downloads of gated models may fail.")
+    logger.warning("HF_TOKEN not found. Downloads of gated models may fail.")
 
 # --- Model Download Function ---
 def download_coz_support_models():
-    print("Checking and downloading CoZ support models...")
+    logger.info("Checking and downloading CoZ support models...")
     for model_key, model_info in CHECKPOINT_FILES_CONFIG.items():
         target_file_path = Path(model_info["target_path"])
         if not target_file_path.exists():
-            print(f"Downloading {model_key} from {model_info['repo_id']}...")
+            logger.info(f"Downloading {model_key} from {model_info['repo_id']}...")
             target_file_path.parent.mkdir(parents=True, exist_ok=True)
             try:
                 cached_file_path = hf_hub_download(
@@ -48,25 +53,48 @@ def download_coz_support_models():
                     filename=model_info['filename'],
                     token=HF_AUTH_TOKEN,
                     cache_dir="./hf_cache",
-                    force_filename=Path(model_info['filename']).name
+                    force_download=False,
+                    resume_download=True
                 )
                 shutil.copy(cached_file_path, target_file_path)
-                print(f"{model_key} downloaded to {target_file_path}")
+                logger.info(f"{model_key} downloaded to {target_file_path}")
             except Exception as e:
-                print(f"Error downloading {model_key}: {e}")
+                logger.error(f"Error downloading {model_key}: {e}")
+                raise
         else:
-            print(f"{model_key} already exists at {target_file_path}")
-    print("All CoZ support models checked.")
+            logger.info(f"{model_key} already exists at {target_file_path}")
+    logger.info("All CoZ support models checked.")
 
 # Download models at startup
-download_coz_support_models()
+try:
+    logger.info("Starting model download...")
+    download_coz_support_models()
+    logger.info("Model download completed.")
+except Exception as e:
+    logger.error(f"Failed to download models: {e}")
+    raise
+
+# --- Preload Stable Diffusion Model ---
+logger.info("Preloading Stable Diffusion model configuration...")
+try:
+    from diffusers import StableDiffusionPipeline
+    StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium-diffusers",
+        use_auth_token=HF_AUTH_TOKEN,
+        cache_dir="./hf_cache"
+    )
+    logger.info("Stable Diffusion model configuration preloaded.")
+except Exception as e:
+    logger.error(f"Failed to preload Stable Diffusion model: {e}")
+    raise
 
 # --- Main Inference Function ---
 def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str, seed: int):
     if input_image is None:
+        logger.error("No input image provided.")
         raise gr.Error("Please upload an image.")
 
-    global HF_AUTH_TOKEN
+    logger.info(f"Starting inference with magnification={magnification}, seed={seed}, caption={caption}")
     with tempfile.TemporaryDirectory() as temp_base_str:
         temp_base_dir = Path(temp_base_str)
         input_img_parent_dir = temp_base_dir / "input_images_root"
@@ -74,12 +102,14 @@ def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str
         input_image_filename = "source_image.png"
         input_image_path = input_img_parent_dir / input_image_filename
         input_image.save(input_image_path, "PNG")
+        logger.info(f"Input image saved to {input_image_path}")
 
         output_img_dir = temp_base_dir / "output_data"
         output_img_dir.mkdir(parents=True, exist_ok=True)
 
         # Check if inference_coz.py exists
         if not Path("inference_coz.py").exists():
+            logger.error("inference_coz.py not found in repository.")
             raise gr.Error("inference_coz.py not found in repository. Please check the Chain-of-Zoom repository.")
 
         command = [
@@ -107,7 +137,7 @@ def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str
         if HF_AUTH_TOKEN:
             command.extend(["--hf_token", HF_AUTH_TOKEN])
 
-        print(f"Running command: {' '.join(command)}")
+        logger.info(f"Running command: {' '.join(command)}")
         process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
 
         stdout_lines = []
@@ -117,7 +147,7 @@ def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str
         # Stream stdout
         if process.stdout:
             for line in iter(process.stdout.readline, ""):
-                print(f"[CoZ STDOUT] {line}", end="")
+                logger.info(f"[CoZ STDOUT] {line.strip()}")
                 stdout_lines.append(line)
                 if "Saving image to" in line:
                     try:
@@ -128,14 +158,14 @@ def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str
         # Stream stderr
         if process.stderr:
             for line in iter(process.stderr.readline, ""):
-                print(f"[CoZ STDERR] {line}", end="")
+                logger.warning(f"[CoZ STDERR] {line.strip()}")
                 stderr_lines.append(line)
 
         process.wait()
 
         if process.returncode != 0:
             error_message = f"Chain-of-Zoom failed.\nSTDOUT:\n{''.join(stdout_lines[-5:])}\nSTDERR:\n{''.join(stderr_lines[-5:])}"
-            print(error_message)
+            logger.error(error_message)
             raise gr.Error(f"Processing failed: {error_message}")
 
         # Find output image
@@ -150,10 +180,11 @@ def run_chain_of_zoom(input_image: Image.Image, magnification: int, caption: str
 
         if not final_output_image_path or not final_output_image_path.exists():
             all_files = list(output_img_dir.rglob("*"))
-            print(f"Files in output directory: {all_files}")
+            logger.error(f"Output image not found in {output_img_dir}. Files found: {all_files}")
             raise gr.Error(f"Output image not found in {output_img_dir}. Files found: {all_files}")
 
         output_image = Image.open(final_output_image_path)
+        logger.info(f"Output image generated: {final_output_image_path}")
         return output_image
 
 # --- Gradio Interface ---
@@ -171,6 +202,7 @@ Optimized for CPU and GPU environments. Ensure HF_TOKEN is set in Space secrets
 """
 article = "<p style='text-align: center;'><a href='https://github.com/bryanswkim/Chain-of-Zoom' target='_blank'>Chain-of-Zoom GitHub</a></p>"
 
+logger.info("Initializing Gradio interface...")
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"<h1 style='text-align: center'>{title}</h1>")
     gr.Markdown(description)
@@ -194,4 +226,10 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    logger.info("Launching Gradio app...")
+    try:
+        demo.launch(server_name="0.0.0.0", server_port=7860)
+        logger.info("Gradio app launched successfully.")
+    except Exception as e:
+        logger.error(f"Failed to launch Gradio app: {e}")
+        raise