added more control for explore endpoint

app.py

@@ -26,8 +26,6 @@ Z_DIM = 100
 DEVICE = torch.device("cpu")
 REPO_ID = "SaniaE/GeoGen"
 FILENAME = "dcgans_model_checkpoint.pt"
-
-# Global model variable
 gen_model = None
 
 @app.on_event("startup")
@@ -49,7 +47,7 @@ def load_model():
         gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
         gen_model.load_state_dict(checkpoint["gen_state_dict"], strict=False)
         gen_model.eval()
-        print("✅ SUCCESS: Petrol Pump GAN is live!")
+        print("SUCCESS: Petrol Pump GAN is live!")
     except Exception as e:
         print(f"Error loading model: {e}")
 
@@ -59,10 +57,8 @@ def postprocess_image(tensor):
     img_tensor = (tensor + 1) / 2
     img_tensor = img_tensor.clamp(0, 1)
 
-    # Use make_grid to handle single or batch images
     grid = vutils.make_grid(img_tensor, padding=0, normalize=False)
     
-    # Convert to HWC format for PIL
     ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
     return Image.fromarray(ndarr)
 
@@ -92,31 +88,32 @@ def generate_random():
     if gen_model is None: raise HTTPException(status_code=503)
     
     with torch.inference_mode():
+        torch.seed() 
+        
         noise = torch.randn(1, Z_DIM, device=DEVICE)
         fake_img = gen_model(noise)
         return StreamingResponse(get_image_stream(fake_img), media_type="image/png")
 
 
 @app.get("/explore")
-def explore_latent(
-    seed: int, 
-    x_shift: float = Query(0.0, ge=-5.0, le=5.0), 
-    y_shift: float = Query(0.0, ge=-5.0, le=5.0)
-):
+def explore_latent(seed: int, x_shift: float = Query(0.0, ge=-5.0, le=5.0), y_shift: float = Query(0.0, ge=-5.0, le=5.0)):
     """Endpoint 2: Controlled generation for 'Tuning'."""
     if gen_model is None: raise HTTPException(status_code=503)
     
     try:
         with torch.inference_mode():
-            # Use the seed to recreate the base 'personality' of the image
             torch.manual_seed(seed)
             if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
-            
+
             noise = torch.randn(1, Z_DIM, device=DEVICE)
 
-            # Apply shifts to specific dimensions
-            noise[0, 0] += x_shift
-            noise[0, 1] += y_shift
+            # Structured control 
+            noise[:, :10] += x_shift
+            noise[:, 10:20] += y_shift
+
+            # Random direction 
+            direction = torch.randn_like(noise)
+            noise = noise + 0.3 * direction * (abs(x_shift) + abs(y_shift))
 
             fake_img = gen_model(noise)
             return StreamingResponse(get_image_stream(fake_img), media_type="image/png")