Update app.py

i fixed the flashing problem earlier the white interface got displayed first before the text could get loaded..

app.py

@@ -5,25 +5,30 @@ from torchvision import models, transforms
 from safetensors.torch import load_file
 from PIL import Image
 
-# 1. Rebuild the blank ResNet-50 architecture
+# 1. Rebuild the skeleton structure globally (Instant execution)
 model = models.resnet50(weights=None)
 model.fc = nn.Linear(model.fc.in_features, 2) # 2 classes: Fake (0) and Real (1)
 
-# 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")
+# Flag to trace if weights have loaded yet
+weights_loaded = False
 
-# (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()
+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!")
 
-# 3. Define the strict patch transform (No randomness!)
+# 2. 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])
@@ -31,11 +36,14 @@ test_transform = transforms.Compose([
 
 classes = ['FAKE', 'REAL']
 
-# 4. Create the prediction function that Gradio will call
+# 3. 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')
     
@@ -43,53 +51,33 @@ def predict_image(img):
     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)}
 
-# 5. Build and launch the Web App!
+# 4. App Copywriting
 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.
-
 """
 
-
-
-# Inside your interface:
-
+# 5. Build and launch the Web App!
 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"
-
 )
 
-interface.launch()
+# Launching with a queue handles the connection states cleanly on Hugging Face
+if __name__ == "__main__":
+    interface.queue().launch()