Update app.py

app.py

@@ -8,6 +8,10 @@ from diffusers import DiffusionPipeline
 import hashlib
 import pickle
 import yaml
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load config file
 with open('config.yaml', 'r') as file:
@@ -16,16 +20,21 @@ with open('config.yaml', 'r') as file:
 # Authenticate using the token stored in Hugging Face Spaces secrets
 if 'HF_TOKEN' in os.environ:
     login(token=os.environ['HF_TOKEN'])
+    logging.info("Successfully logged in with HF_TOKEN")
 else:
-    raise ValueError("HF_TOKEN not found in environment variables. Please add it to your Space's secrets.")
+    logging.warning("HF_TOKEN not found in environment variables. Some functionality may be limited.")
 
 # Correctly access the config values
 process_config = config['config']['process'][0]  # Assuming the first process is the one we want
 
-base_model = process_config['model']['name_or_path']
+base_model = "black-forest-labs/FLUX.1-dev"
 lora_model = "sagar007/sagar_flux"  # This isn't in the config, so we're keeping it as is
 trigger_word = process_config['trigger_word']
 
+logging.info(f"Base model: {base_model}")
+logging.info(f"LoRA model: {lora_model}")
+logging.info(f"Trigger word: {trigger_word}")
+
 # Global variables
 pipe = None
 cache = {}
@@ -46,15 +55,13 @@ def initialize_model():
     global pipe
     if pipe is None:
         try:
-            print("Loading base model...")
-            pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16)
-            print("Moving model to CUDA...")
+            logging.info(f"Attempting to load model: {base_model}")
+            pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.float16, use_safetensors=True)
+            logging.info("Moving model to CUDA...")
             pipe = pipe.to("cuda")
-            print(f"Successfully loaded base model: {base_model}")
+            logging.info(f"Successfully loaded model: {base_model}")
         except Exception as e:
-            print(f"Error initializing model: {str(e)}")
-            import traceback
-            print(traceback.format_exc())
+            logging.error(f"Error loading model {base_model}: {str(e)}")
             raise
 
 def load_cache():
@@ -62,12 +69,12 @@ def load_cache():
     if os.path.exists(CACHE_FILE):
         with open(CACHE_FILE, 'rb') as f:
             cache = pickle.load(f)
-    print(f"Loaded {len(cache)} cached images")
+    logging.info(f"Loaded {len(cache)} cached images")
 
 def save_cache():
     with open(CACHE_FILE, 'wb') as f:
         pickle.dump(cache, f)
-    print(f"Saved {len(cache)} cached images")
+    logging.info(f"Saved {len(cache)} cached images")
 
 def get_cache_key(prompt, cfg_scale, steps, seed, width, height, lora_scale):
     return hashlib.md5(f"{prompt}{cfg_scale}{steps}{seed}{width}{height}{lora_scale}".encode()).hexdigest()
@@ -82,23 +89,23 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
     cache_key = get_cache_key(prompt, cfg_scale, steps, seed, width, height, lora_scale)
 
     if cache_key in cache:
-        print("Using cached image")
+        logging.info("Using cached image")
         return cache[cache_key], seed
 
     try:
-        print(f"Starting run_lora with prompt: {prompt}")
+        logging.info(f"Starting run_lora with prompt: {prompt}")
         if pipe is None:
-            print("Initializing model...")
+            logging.info("Initializing model...")
             initialize_model()
 
-        print(f"Using seed: {seed}")
+        logging.info(f"Using seed: {seed}")
 
         generator = torch.Generator(device="cuda").manual_seed(seed)
 
         full_prompt = f"{prompt} {trigger_word}"
-        print(f"Full prompt: {full_prompt}")
+        logging.info(f"Full prompt: {full_prompt}")
 
-        print("Starting image generation...")
+        logging.info("Starting image generation...")
         image = pipe(
             prompt=full_prompt,
             num_inference_steps=steps,
@@ -107,7 +114,7 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
             height=height,
             generator=generator,
         ).images[0]
-        print("Image generation completed successfully")
+        logging.info("Image generation completed successfully")
 
         # Cache the generated image
         cache[cache_key] = image
@@ -115,9 +122,9 @@ def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora
 
         return image, seed
     except Exception as e:
-        print(f"Error during generation: {str(e)}")
+        logging.error(f"Error during generation: {str(e)}")
         import traceback
-        print(traceback.format_exc())
+        logging.error(traceback.format_exc())
         return None, seed
 
 def update_prompt(example):
@@ -164,8 +171,8 @@ with gr.Blocks() as app:
 
 # Launch the app
 if __name__ == "__main__":
-    print("Starting the Gradio app...")
-    print("Pre-generating example images...")
+    logging.info("Starting the Gradio app...")
+    logging.info("Pre-generating example images...")
     cache_example_images()
     app.launch(share=True)
-    print("Gradio app launched successfully")
+    logging.info("Gradio app launched successfully")