Update app.py

app.py

@@ -16,6 +16,8 @@ load_dotenv()
 
 # Retrieve Hugging Face token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("Hugging Face token not found in environment variables.")
 
 ## 2.1 Image Analysis with DETR
 def load_detr_model():
@@ -24,7 +26,7 @@ def load_detr_model():
         detr_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
         return detr_model, detr_processor, None
     except Exception as e:
-        return None, None, f"Error loading DETR model: {e}"
+        return None, None, f"Error loading DETR model: {str(e)}"
 
 detr_model, detr_processor, detr_error = load_detr_model()
 
@@ -37,7 +39,7 @@ def detect_objects(image):
             results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
             return results, None
         except Exception as e:
-            return None, f"Error in detect_objects: {e}"
+            return None, f"Error in detect_objects: {str(e)}"
     else:
         return None, "DETR models not loaded. Skipping object detection."
 
@@ -76,7 +78,7 @@ def style_transfer(content_image, style_image):
         generated_image = generated.squeeze().clamp(0, 255).cpu().detach().numpy().transpose(1, 2, 0)
         return Image.fromarray(np.uint8(generated_image)), None
     except Exception as e:
-        return content_image, f"Error in style_transfer: {e}"
+        return content_image, f"Error in style_transfer: {str(e)}"
 
 ## 2.3 Layout Generation with LayoutLM
 def load_layoutlm_model():
@@ -85,7 +87,7 @@ def load_layoutlm_model():
         layoutlm_model = LayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased')
         return layoutlm_tokenizer, layoutlm_model, None
     except Exception as e:
-        return None, None, f"Error loading LayoutLM model: {e}"
+        return None, None, f"Error loading LayoutLM model: {str(e)}"
 
 layoutlm_tokenizer, layoutlm_model, layoutlm_error = load_layoutlm_model()
 
@@ -97,20 +99,18 @@ def generate_layout(text):
             layout = outputs.logits.argmax(dim=-1)
             return layout, None
         except Exception as e:
-            return None, f"Error in generate_layout: {e}"
+            return None, f"Error in generate_layout: {str(e)}"
     else:
         return None, "LayoutLM models not loaded. Skipping layout generation."
 
 ## 2.4 Image Generation with Stable Diffusion
 def load_stable_diffusion_model():
     try:
-        if HF_TOKEN is None:
-            raise ValueError("Hugging Face token not found in environment variables.")
         login(token=HF_TOKEN)
         sd_pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cuda")
         return sd_pipeline, None
     except Exception as e:
-        return None, f"Error loading Stable Diffusion model: {e}"
+        return None, f"Error loading Stable Diffusion model: {str(e)}"
 
 sd_pipeline, sd_error = load_stable_diffusion_model()
 
@@ -120,7 +120,7 @@ def generate_image(prompt):
             image = sd_pipeline(prompt).images[0]
             return image, None
         except Exception as e:
-            return None, f"Error in generate_image: {e}"
+            return None, f"Error in generate_image: {str(e)}"
     else:
         return None, "Stable Diffusion model not loaded. Skipping image generation."
 
@@ -130,7 +130,7 @@ def load_upscale_pipeline():
         upscale_pipeline = StableDiffusionUpscalePipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler").to("cuda")
         return upscale_pipeline, None
     except Exception as e:
-        return None, f"Error loading Upscale Pipeline: {e}"
+        return None, f"Error loading Upscale Pipeline: {str(e)}"
 
 upscale_pipeline, upscale_error = load_upscale_pipeline()
 
@@ -142,7 +142,7 @@ def super_resolve(image):
             upscaled_image = upscale_pipeline(image=image).images[0]
             return upscaled_image, None
         except Exception as e:
-            return None, f"Error in super_resolve: {e}"
+            return None, f"Error in super_resolve: {str(e)}"
     else:
         return image, "Upscale Pipeline not loaded. Skipping super-resolution."
 
@@ -176,7 +176,7 @@ def process_image(image, style_image, text_prompt):
 
         return final_image, None
     except Exception as e:
-        return None, f"Error in process_image: {e}"
+        return None, f"Error in process_image: {str(e)}"
 
 iface = gr.Interface(
     fn=process_image,
@@ -194,5 +194,5 @@ iface = gr.Interface(
 try:
     iface.launch()
 except Exception as e:
-    print(f"Error occurred while launching the interface: {e}")
+    print(f"Error occurred while launching the interface: {str(e)}")
     traceback.print_exc()