Add logging and static route for image serving

app.py

@@ -12,7 +12,7 @@ logging.basicConfig(level=logging.DEBUG)
 
 # Define paths
 MODEL_PATH = "Roshan1162003/fine_tuned_model"
-STATIC_IMAGES_PATH = os.path.join("/tmp", "images")
+STATIC_IMAGES_PATH = os.path.join("static", "images")
 os.makedirs(STATIC_IMAGES_PATH, exist_ok=True)
 logging.info(f"Created static directory: {STATIC_IMAGES_PATH}")
 
@@ -22,28 +22,26 @@ RESTRICTED_TERMS = [
     "offensive", "hate", "nude", "porn", "gore", "drug"
 ]
 
-# Verify HF_TOKEN
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    logging.error("HF_TOKEN environment variable is not set")
-    raise EnvironmentError("HF_TOKEN environment variable is not set")
-
 # Load the fine-tuned model
 pipe = None
 try:
-    logging.info(f"Loading model from {MODEL_PATH} with HF_TOKEN")
-    dtype = torch.float32
-    logging.info(f"Using dtype: {dtype}")
-    pipe = StableDiffusionPipeline.from_pretrained(
-        MODEL_PATH,
-        torch_dtype=dtype,
-        use_safetensors=True,
-        use_auth_token=HF_TOKEN
-    ).to("cpu")
-    logging.info("Model loaded successfully with default scheduler")
+    if torch.cuda.is_available():
+        pipe = StableDiffusionPipeline.from_pretrained(
+            MODEL_PATH,
+            torch_dtype=torch.float16,
+            use_safetensors=True,
+            use_auth_token=os.getenv("HF_TOKEN")
+        ).to("cuda")
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained(
+            MODEL_PATH,
+            torch_dtype=torch.float32,
+            use_safetensors=True,
+            use_auth_token=os.getenv("HF_TOKEN")
+        )
+    logging.info("Model loaded successfully")
 except Exception as e:
-    logging.error(f"Failed to load model: {str(e)}")
-    raise RuntimeError(f"Failed to load model: {str(e)}")
+    logging.error(f"Error loading model: {e}")
 
 # Aspect ratio to resolution mapping
 ASPECT_RATIOS = {
@@ -54,7 +52,7 @@ ASPECT_RATIOS = {
 
 def is_prompt_safe(prompt):
     """Check if prompt contains restricted terms."""
-    prompt_lower = prompt.lower() if isinstance(prompt, str) else ""
+    prompt_lower = prompt.lower()
     for term in RESTRICTED_TERMS:
         if re.search(r'\b' + re.escape(term) + r'\b', prompt_lower):
             return False
@@ -68,70 +66,43 @@ def index():
 @app.route("/generate", methods=["POST"])
 def generate():
     try:
-        # Log full request data
-        logging.info(f"Request headers: {request.headers}")
-        logging.info(f"Request form data: {request.form}")
-
         prompt = request.form.get("prompt", "").strip()
-        num_images = request.form.get("num_images")
+        num_images = int(request.form.get("num_images", 1))
         aspect_ratio = request.form.get("aspect_ratio", "1:1")
         model_name = request.form.get("model", "stable_diffusion")
+        logging.info(f"Received prompt: {prompt}, num_images: {num_images}, aspect_ratio: {aspect_ratio}")
 
-        # Validate inputs
-        if not prompt or not isinstance(prompt, str):
-            logging.error("Invalid prompt: empty or not a string")
-            return jsonify({"error": "Prompt is required and must be a string"}), 400
-        try:
-            num_images = int(num_images) if num_images else 1
-        except (ValueError, TypeError):
-            logging.error(f"Invalid num_images: {num_images}")
-            return jsonify({"error": "Number of images must be an integer"}), 400
+        if not prompt:
+            return jsonify({"error": "Prompt is required"}), 400
+        if model_name != "stable_diffusion":
+            return jsonify({"error": "Selected model is locked"}), 400
         if num_images < 1 or num_images > 5:
-            logging.error(f"Invalid num_images: {num_images}")
             return jsonify({"error": "Number of images must be between 1 and 5"}), 400
         if aspect_ratio not in ASPECT_RATIOS:
-            logging.error(f"Invalid aspect_ratio: {aspect_ratio}")
             return jsonify({"error": "Invalid aspect ratio"}), 400
-        if model_name != "stable_diffusion":
-            logging.error(f"Invalid model: {model_name}")
-            return jsonify({"error": "Selected model is locked"}), 400
-
-        logging.info(f"Received prompt: {prompt}, num_images: {num_images}, aspect_ratio: {aspect_ratio}")
 
         if not is_prompt_safe(prompt):
-            logging.error(f"Restricted prompt detected: {prompt}")
             return jsonify({
                 "error": "You are violating the regulation policy terms and conditions due to restricted terms in the prompt."
             }), 400
 
-        if pipe is None:
-            logging.error("Pipeline is None, cannot generate images")
-            return jsonify({"error": "Model pipeline is not initialized"}), 500
-
         width, height = ASPECT_RATIOS[aspect_ratio]
         image_paths = []
-        negative_prompt = "blurry, low quality, distorted face, unnatural features, overexposed, underexposed"
         for i in range(num_images):
-            logging.info(f"Generating image {i+1}/{num_images} for prompt: {prompt}")
-            start_time = time.time()
             image = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
+                prompt,
                 width=width,
                 height=height,
                 num_inference_steps=50,
                 guidance_scale=7.5,
                 seed=42 + i
             ).images[0]
-            end_time = time.time()
-            logging.info(f"Image generation took {end_time - start_time:.2f} seconds")
             timestamp = int(time.time() * 1000)
             image_path = os.path.join(STATIC_IMAGES_PATH, f"generated_{timestamp}_{i}.png")
             logging.info(f"Saving image to: {image_path}")
             image.save(image_path)
             logging.info(f"Image saved successfully: {image_path}")
-            image_paths.append(image_path.replace("/tmp/", ""))
-            logging.info(f"Image path for UI: {image_paths[-1]}")
+            image_paths.append(image_path.replace("static/", ""))
         logging.info(f"Returning image paths: {image_paths}")
         return jsonify({"images": image_paths})