Update app.py

app.py

@@ -1,9 +1,8 @@
 import os
 import torch
-import legacy
 import numpy as np
 import PIL.Image
-from flask import Flask, request, render_template, send_file
+from flask import Flask, request, render_template, send_file, url_for
 from werkzeug.utils import secure_filename
 
 # Set up Flask app
@@ -18,60 +17,96 @@ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(RESULT_FOLDER, exist_ok=True)
 
 # Load StyleGAN3 Model
-network_pkl = "models/stylegan3-r-ffhq-1024x1024.pkl"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print(f'Loading networks from "{network_pkl}"...')
-with open(network_pkl, "rb") as f:
-    G = legacy.load_network_pkl(f)["G_ema"].to(device)
+try:
+    import legacy
+    
+    network_pkl = "models/stylegan3-r-ffhq-1024x1024.pkl"
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f'Loading networks from "{network_pkl}"...')
+    
+    with open(network_pkl, "rb") as f:
+        G = legacy.load_network_pkl(f)["G_ema"].to(device)
+    
+    model_loaded = True
+    print("StyleGAN3 model loaded successfully!")
+except Exception as e:
+    print(f"Error loading StyleGAN3 model: {e}")
+    model_loaded = False
 
 # Function to encode an image into latent space
 def image_to_latent(image_path):
+    # Note: This is a simplified version. Actual image encoding to latent space
+    # requires techniques like optimization or encoder networks
     latent_vector = torch.randn(1, G.z_dim, device=device)  # Generate random latent vector
     return latent_vector
 
 # Function to modify latent code to make the face look younger
 def modify_age(latent_vector, age_factor=-2.0):
-    age_direction = torch.load("models/age_direction.pt").to(device)  # Load precomputed age direction
-    new_latent_vector = latent_vector + age_factor * age_direction
-    return new_latent_vector
+    try:
+        age_direction = torch.load("models/age_direction.pt").to(device)  # Load precomputed age direction
+        new_latent_vector = latent_vector + age_factor * age_direction
+        return new_latent_vector
+    except Exception as e:
+        print(f"Error modifying age: {e}")
+        return latent_vector  # Return original if error
 
 # Function to generate an image from a latent code
 def generate_image(latent_vector):
-    img = G.synthesis(latent_vector, noise_mode="const")
-    img = (img + 1) * (255 / 2)
-    img = img.permute(0, 2, 3, 1).cpu().numpy()[0].astype(np.uint8)
-    return PIL.Image.fromarray(img)
+    try:
+        img = G.synthesis(latent_vector, noise_mode="const")
+        img = (img + 1) * (255 / 2)
+        img = img.permute(0, 2, 3, 1).cpu().numpy()[0].astype(np.uint8)
+        return PIL.Image.fromarray(img)
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        # Return a blank image if there's an error
+        return PIL.Image.new('RGB', (1024, 1024), color = (255, 255, 255))
 
 # Flask Routes
 @app.route("/", methods=["GET", "POST"])
 def upload_file():
+    error_message = None
+    
+    if not model_loaded:
+        error_message = "StyleGAN3 model could not be loaded. Please check the server logs."
+    
     if request.method == "POST":
         if "file" not in request.files:
-            return "No file uploaded", 400
+            return render_template("index.html", error="No file uploaded")
         
         file = request.files["file"]
         if file.filename == "":
-            return "No selected file", 400
-
-        filename = secure_filename(file.filename)
-        input_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
-        file.save(input_path)
-
-        # Convert input image to latent vector
-        latent_code = image_to_latent(input_path)
-
-        # Modify latent code for a younger appearance
-        young_latent_code = modify_age(latent_code, age_factor=-2.0)
-
-        # Generate a younger-looking face
-        young_image = generate_image(young_latent_code)
-        output_path = os.path.join(app.config["RESULT_FOLDER"], "young_" + filename)
-        young_image.save(output_path)
-
-        return render_template("result.html", filename="young_" + filename)
-
-    return render_template("index.html")
+            return render_template("index.html", error="No selected file")
+        
+        if not model_loaded:
+            return render_template("index.html", error=error_message)
+        
+        try:
+            filename = secure_filename(file.filename)
+            input_path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
+            file.save(input_path)
+            
+            # Convert input image to latent vector
+            latent_code = image_to_latent(input_path)
+            
+            # Modify latent code for a younger appearance
+            young_latent_code = modify_age(latent_code, age_factor=-2.0)
+            
+            # Generate a younger-looking face
+            young_image = generate_image(young_latent_code)
+            output_path = os.path.join(app.config["RESULT_FOLDER"], "young_" + filename)
+            young_image.save(output_path)
+            
+            return render_template("result.html", filename="young_" + filename)
+        except Exception as e:
+            error_message = f"Error processing image: {str(e)}"
+            return render_template("index.html", error=error_message)
+    
+    return render_template("index.html", error=error_message)
+
+@app.route("/results/<filename>")
+def display_image(filename):
+    return send_file(os.path.join(app.config["RESULT_FOLDER"], filename))
 
 @app.route("/download/<filename>")
 def download_file(filename):
@@ -79,4 +114,4 @@ def download_file(filename):
 
 # Run the Flask app
 if __name__ == "__main__":
-    app.run(debug=True)
+    app.run(debug=True, host="0.0.0.0")