created fastapi

app.py

@@ -0,0 +1,118 @@
+import os
+import io
+import torch
+from torch import nn 
+from PIL import Image
+import torchvision.utils as vutils
+from fastapi import FastAPI, Response, HTTPException, Query
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import hf_hub_download
+
+from models import Generator
+
+app = FastAPI()
+
+# CORS Configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Configuration Constants
+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")
+def load_model():
+    global gen_model
+    try:
+        token = os.getenv("HF_TOKEN")
+        model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, token=token)
+        
+        checkpoint = torch.load(model_path, map_location=DEVICE)
+        
+        gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
+        gen_model.load_state_dict(checkpoint["gen_state_dict"])
+        gen_model.eval()
+        print("Model loaded successfully.")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+
+
+def postprocess_image(tensor):
+    # Unnormalize: tanh output [-1, 1] -> [0, 1]
+    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)
+
+
+def get_image_stream(tensor):
+    """Helper to convert tensor to a streaming-ready PNG."""
+    img_tensor = (tensor + 1) / 2
+    img_tensor = img_tensor.clamp(0, 1)
+    grid = vutils.make_grid(img_tensor, padding=0)
+    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    
+    pil_img = Image.fromarray(ndarr)
+    buf = io.BytesIO()
+    pil_img.save(buf, format="PNG")
+    buf.seek(0)
+    return buf
+
+
+@app.get("/")
+def read_root():
+    return {"status": "online", "model": REPO_ID}
+
+
+@app.get("/generate")
+def generate_random():
+    """Endpoint 1: Purely random generation for 'Discovery'."""
+    if gen_model is None: raise HTTPException(status_code=503)
+    
+    with torch.inference_mode():
+        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)
+):
+    """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
+
+            fake_img = gen_model(noise)
+            return StreamingResponse(get_image_stream(fake_img), media_type="image/png")
+            
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

models.py

@@ -0,0 +1,31 @@
+from torch import nn 
+
+class Generator(nn.Module):
+    def __init__(self, z_dim=100, input_channels=3, hidden_dim=64):
+        super(Generator, self).__init__()
+        self.z_dim = z_dim
+        self.gen = nn.Sequential(
+            self.generator_block(z_dim, hidden_dim * 32, stride=1, padding=0),  
+            self.generator_block(hidden_dim * 32, hidden_dim * 16), 
+            self.generator_block(hidden_dim * 16, hidden_dim * 8),  
+            self.generator_block(hidden_dim * 8, hidden_dim * 4),  
+            self.generator_block(hidden_dim * 4, hidden_dim * 2),  
+            self.generator_block(hidden_dim * 2, hidden_dim), 
+            self.generator_block(hidden_dim, input_channels, final_layer=True) 
+        )
+        
+    def generator_block(self, input_channels, output_channels, kernel_size=4, stride=2, padding=1, final_layer=False): 
+        if not final_layer:
+            return nn.Sequential(
+                nn.ConvTranspose2d(input_channels, output_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
+                nn.InstanceNorm2d(output_channels, affine=True),  
+                nn.ReLU(inplace=True)
+            )
+        else: 
+            return nn.Sequential(
+                nn.ConvTranspose2d(input_channels, output_channels, kernel_size=kernel_size, stride=stride, padding=padding),
+                nn.Tanh()
+            )
+
+    def forward(self, noise):
+        return self.gen(noise.view(len(noise), self.z_dim, 1, 1))