Update app.py

app.py

@@ -1,8 +1,10 @@
 import os
 import subprocess
 
+# Install Gradio
 os.system("pip install gradio==3.50")
 
+# Create and run setup script for CUDA environment variables
 with open('setup.sh', 'w') as f:
     f.write('''#!/bin/bash
 export CUDA_HOME=/usr/local/cuda
@@ -19,19 +21,19 @@ os.system('./setup.sh')
 import torch
 
 # Check CUDA availability
-print("CUDA available:", torch.cuda.is_available())
-
+is_cuda_available = torch.cuda.is_available()
+device = 'cuda' if is_cuda_available else 'cpu'
+print("CUDA available:", is_cuda_available)
 
 from argparse import Namespace
 import pprint
 import numpy as np
 from PIL import Image
-import torch
 import torchvision.transforms as transforms
 import cv2
-import dlibs.dlib
+import dlib
 import matplotlib.pyplot as plt
-import gradio as gr  # Importing Gradio as gr
+import gradio as gr
 from tensorflow.keras.preprocessing.image import img_to_array
 from huggingface_hub import hf_hub_download, login
 from datasets.augmentations import AgeTransformer
@@ -67,14 +69,14 @@ EXPERIMENT_DATA_ARGS = {
 }
 EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[EXPERIMENT_TYPE]
 model_path = EXPERIMENT_ARGS['model_path']
-ckpt = torch.load(model_path, map_location='cpu')
+ckpt = torch.load(model_path, map_location=device)
 opts = ckpt['opts']
 pprint.pprint(opts)
 opts['checkpoint_path'] = model_path
 opts = Namespace(**opts)
 net = pSp(opts)
 net.eval()
-net.cuda()
+net.to(device)  # Send model to the appropriate device (GPU or CPU)
 
 print('Model successfully loaded!')
 
@@ -149,9 +151,9 @@ def apply_aging(image, target_age):
     results = []
     for age_transformer in age_transformers:
         with torch.no_grad():
-            input_image_age = [age_transformer(input_image.cpu()).to('cuda')]
+            input_image_age = [age_transformer(input_image.cpu()).to(device)]  # Move to appropriate device
             input_image_age = torch.stack(input_image_age)
-            result_tensor = net(input_image_age.to("cuda").float(), randomize_noise=False, resize=False)[0]
+            result_tensor = net(input_image_age.to(device).float(), randomize_noise=False, resize=False)[0]
             result_image = tensor2im(result_tensor)
             results.append(np.array(result_image))
     final_result = results[0]
@@ -184,10 +186,13 @@ def process_image(uploaded_image):
 
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
-    title="Aging Effect with Teeth Replacement",
-    description="Upload an image to apply an aging effect. The application will generate two results: one with good teeth and one with bad teeth."
+    inputs=gr.Image(type="pil", label="Upload an Image"),
+    outputs=[
+        gr.Image(label="Image with Good Teeth"),
+        gr.Image(label="Image with Bad Teeth")
+    ],
+    title="Age Transformation and Teeth Replacement",
+    description="Upload an image and apply aging effects while replacing teeth."
 )
 
-iface.launch(debug=True)
+iface.launch(share=True)