Update app.py

app.py

@@ -1,39 +1,19 @@
 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
-export PATH=$CUDA_HOME/bin:$PATH
-export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
-''')
-
-# Make the script executable
-os.system('chmod +x setup.sh')
-
-# Run the setup script to set environment variables
-os.system('./setup.sh')
-
-import torch
-
-# Check CUDA availability
-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 matplotlib.pyplot as plt
-import gradio as gr
+import gradio as gr  # Importing 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
@@ -69,14 +49,14 @@ EXPERIMENT_DATA_ARGS = {
 }
 EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[EXPERIMENT_TYPE]
 model_path = EXPERIMENT_ARGS['model_path']
-ckpt = torch.load(model_path, map_location=device)
+ckpt = torch.load(model_path, map_location='cpu')
 opts = ckpt['opts']
 pprint.pprint(opts)
 opts['checkpoint_path'] = model_path
 opts = Namespace(**opts)
 net = pSp(opts)
 net.eval()
-net.to(device)  # Send model to the appropriate device (GPU or CPU)
+net.cuda()
 
 print('Model successfully loaded!')
 
@@ -151,9 +131,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(device)]  # Move to appropriate device
+            input_image_age = [age_transformer(input_image.cpu()).to('cuda')]
             input_image_age = torch.stack(input_image_age)
-            result_tensor = net(input_image_age.to(device).float(), randomize_noise=False, resize=False)[0]
+            result_tensor = net(input_image_age.to("cuda").float(), randomize_noise=False, resize=False)[0]
             result_image = tensor2im(result_tensor)
             results.append(np.array(result_image))
     final_result = results[0]
@@ -186,13 +166,10 @@ def process_image(uploaded_image):
 
 iface = gr.Interface(
     fn=process_image,
-    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."
+    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."
 )
 
-iface.launch(share=True)
+iface.launch(debug=True)