Update app.py

reverted back to my previous version

app.py

@@ -1,83 +1,189 @@
 import gradio as gr
+
 import torch
+
 import torch.nn as nn
+
 from torchvision import models, transforms
+
 from safetensors.torch import load_file
+
 from PIL import Image
 
-# 1. Rebuild the skeleton structure globally (Instant execution)
+
+
+# 1. Rebuild the blank ResNet-50 architecture
+
 model = models.resnet50(weights=None)
+
 model.fc = nn.Linear(model.fc.in_features, 2) # 2 classes: Fake (0) and Real (1)
 
-# Flag to trace if weights have loaded yet
-weights_loaded = False
 
-def load_model_weights():
-    global weights_loaded, model
-    if not weights_loaded:
-        print(" Loading model weights into memory...")
-        # Load weights from the safetensors file on CPU
-        state_dict = load_file("model.safetensors", device="cpu")
-        
-        # Clean DataParallel prefix strings if they exist
-        state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
-        
-        model.load_state_dict(state_dict)
-        model.eval()
-        weights_loaded = True
-        print(" Weights loaded successfully!")
 
-# 2. Define the strict patch transform (No randomness!)
+# 2. Load YOUR trained weights from the safetensors file
+
+# Using map_location='cpu' ensures it works on Hugging Face's free CPU tier!
+
+state_dict = load_file("model.safetensors", device="cpu")
+
+
+
+# (Optional safety check) If you trained with DataParallel, keys might have "module." in front.
+
+# This removes it so the weights load perfectly no matter what.
+
+state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+
+
+
+# Inject the weights into the skeleton and set to test mode
+
+model.load_state_dict(state_dict)
+
+model.eval()
+
+
+
+# 3. Define the strict patch transform (No randomness!)
+
 test_transform = transforms.Compose([
+
+    # transforms.Resize(256), # Standard practice to resize slightly before center cropping
+
     transforms.CenterCrop(224), 
+
     transforms.ToTensor(),
+
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+
 ])
 
+
+
 classes = ['FAKE', 'REAL']
 
-# 3. Create the prediction function that Gradio will call
+
+
+# 4. Create the prediction function that Gradio will call
+
 def predict_image(img):
+
     if img is None:
+
         return None
+
         
-    # Lazy-load weights on the first click to prevent UI setup blocking
-    load_model_weights()
-        
+
     # Ensure image is strictly RGB (drops alpha channels from PNGs)
+
     img = img.convert('RGB')
+
     
+
     # Apply transforms and add the batch dimension
+
     img_tensor = test_transform(img).unsqueeze(0) 
+
     
+
     with torch.no_grad():
+
+        # Pass through model
+
         preds = model(img_tensor)
+
+        # Convert raw numbers to percentages (0.0 to 1.0)
+
         probs = torch.nn.functional.softmax(preds[0], dim=0)
+
         
+
+    # Gradio expects a dictionary of { "Class Name": probability }
+
     return {classes[i]: float(probs[i]) for i in range(2)}
 
-# 4. App Copywriting
+
+
+# 5. Build and launch the Web App!
+
 description_text = """
+
+
+
 ###  How it works:
+
+
+
 This model analyzes microscopic pixel noise to determine if an image is real or AI-generated. 
 
+
+
+
+
+
+
 ###  Limitations for best results:
+
+
+
 * **Resolution Sweet Spot:** Works flawlessly on standard AI resolutions and mid-sized images (from **512x512 up to around 1000x1500 pixels**, like 640x832 or 880x1320). 
+
+
+
 * **The 4K Danger Zone:** Ultra-high-resolution (like **3840x2160 / 4K**) images will cause the model to fail. Because the model's 'magnifying glass' is strictly fixed to a 224x224 pixel crop, it ends up looking through a pinhole at less than 0.6% of a 4K image, causing it to lose context and guess randomly.
+
+
+
 * **Centered Subjects:** The model strictly scans the dead-center of the image. If the AI artifacts or mistakes (like extra fingers or warped backgrounds) are on the far edges, the model won't see them!
+
+
+
 * **No Screenshots:** Heavy compression (like taking a screenshot or downloading from messaging apps) destroys the microscopic forensic evidence. Please upload the raw, original files.
+
+
+
 """
 
-# 5. Build and launch the Web App!
+
+
+
+
+
+
+# Inside your interface:
+
+
+
 interface = gr.Interface(
+
+
+
     fn=predict_image,
+
+
+
     inputs=gr.Image(type="pil", label="Upload an Image"),
+
+
+
     outputs=gr.Label(num_top_classes=2, label="Prediction"),
+
+
+
     title="PixelSleuth: AI Image Detector",
+
+
+
     description=description_text, 
+
+
+
     flagging_mode="never"
+
+
+
 )
 
-# Launching with a queue handles the connection states cleanly on Hugging Face
-if __name__ == "__main__":
-    interface.queue().launch()
+
+
+interface.launch()