Update app.py

app.py

@@ -8,9 +8,9 @@ from transformers import DetrImageProcessor, DetrForObjectDetection
 from diffusers import StableDiffusionPipeline
 from huggingface_hub import login
 import torchvision.transforms as T
+from dotenv import load_dotenv
 
 # Load environment variables from .env file
-from dotenv import load_dotenv
 load_dotenv()
 
 # Retrieve Hugging Face token from environment variable
@@ -18,6 +18,9 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
     raise ValueError("Hugging Face token not found in environment variables.")
 
+# Login to Hugging Face using the token
+login(token=HF_TOKEN)
+
 # Load DETR model for object detection
 def load_detr_model():
     try:
@@ -45,8 +48,8 @@ def detect_objects(image):
 # Load Stable Diffusion model for image generation
 def load_stable_diffusion_model():
     try:
-        login(token=HF_TOKEN)
-        pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cuda")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
         return pipeline, None
     except Exception as e:
         return None, f"Error loading Stable Diffusion model: {str(e)}"
@@ -72,7 +75,6 @@ def process_image(image):
             return None, detect_error
 
         # Generate a modern redesign of the image based on the detected objects
-        # For simplicity, we'll use a fixed prompt for image generation
         prompt = "modern redesign of an interior room"
         generated_image, gen_image_error = generate_image(prompt)
         if gen_image_error:
@@ -97,4 +99,4 @@ try:
     iface.launch()
 except Exception as e:
     print(f"Error occurred while launching the interface: {str(e)}")
-    traceback.print_exc()
+    traceback.print_exc()