Update app.py

app.py

@@ -5,8 +5,8 @@ from PIL import Image
 from torchvision import transforms
 import random
 from huggingface_hub import hf_hub_download
-from generator_1 import Generator as Generator_1  # Import your StyleGAN2 generator
-from generator_2 import Generator as Generator_2  # Import your SRGAN generator
+from generator_1 import Generator as StyleGANGenerator  # Import your StyleGAN2 generator
+from generator_2 import Generator as SRGANGenerator  # Import your SRGAN generator
 
 # wts = ['trial_0_G (1).pth', 'trial_0_G (2).pth', 'trial_0_G (3).pth', 'trial_0_G (4).pth', 'trial_0_G (5).pth', 'trial_0_G.pth']
 wts = ['trial_0_G (2).pth', 'trial_0_G (5).pth', 'trial_0_G.pth']
@@ -14,8 +14,8 @@ random_wt = random.choice(wts)
 
 # Load trained model weights from Hugging Face Hub
 weights_path = hf_hub_download(
-    repo_id="keysun89/image_generation",  # Replace with your repo
-    filename=random_wt  # Replace with your weights file
+    repo_id="keysun89/img_generation",  # Fixed repo name
+    filename=random_wt
 )
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -26,7 +26,7 @@ w_dim = 512
 img_resolution = 256  # Adjust to your training resolution
 img_channels = 3
 
-model = Generator_1(
+model = StyleGANGenerator(
     z_dim=z_dim,
     w_dim=w_dim,
     img_resolution=img_resolution,
@@ -38,13 +38,14 @@ model.load_state_dict(torch.load(weights_path, map_location=device))
 model.to(device)
 model.eval()
 
+wt_2 = 'genrator.pth'
 srgan_weights = hf_hub_download(
-    repo_id="keysun89/image_generation",
-    filename= "genrator.pth"
+    repo_id="keysun89/img_generation",  # Fixed repo name
+    filename=wt_2
 )
 
 # Initialize SRGAN with scale=2 (256 -> 512)
-srgan_model = Generator_2(img_feat=3, n_feats=64, kernel_size=3, num_block=16, scale=2)
+srgan_model = SRGANGenerator(img_feat=3, n_feats=64, kernel_size=3, num_block=16, scale=2)
 srgan_model.load_state_dict(torch.load(srgan_weights, map_location=device))
 srgan_model.to(device)
 srgan_model.eval()
@@ -70,14 +71,14 @@ def generate():
         # Step 2: Upscale to 512x512 with SRGAN
         img_256_tensor = transform(img_256_pil).unsqueeze(0).to(device)
         
-        # Generate high-resolution image
-        img_512 = srgan_model(img_256_tensor)
+        # Generate high-resolution image (SRGAN returns tuple: image, features)
+        img_512, _ = srgan_model(img_256_tensor)
         
         # Convert to PIL Image (512x512)
         img_512_np = img_512.squeeze(0).cpu().numpy()
         img_512_np = np.transpose(img_512_np, (1, 2, 0))  # CHW to HWC
-        # Denormalize if needed
-        img_512_np = (img_512_np * 127.5 + 128).clip(0, 255).astype(np.uint8)
+        # Denormalize from tanh output [-1, 1] to [0, 255]
+        img_512_np = (img_512_np * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
         img_512_pil = Image.fromarray(img_512_np)
     
     return img_256_pil, img_512_pil