Changed network snapshot

app.py

@@ -13,13 +13,18 @@ import os
 
 network_pkl_a = 'stylegan3-r-afhqv2-512x512.pkl'
 network_pkl_d = 'network-snapshot.pkl'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Device = ",device)
+
 with open(network_pkl_d, 'rb') as f:
-    G_d = pickle.load(f)['G_ema'].cpu()  # torch.nn.Module
+    G_d = pickle.load(f)['G_ema'].to(device) #.cpu()  # torch.nn.Module
 
 with open(network_pkl_a, 'rb') as f:
-    G_a = pickle.load(f)['G_ema'].cpu()  # torch.nn.Module
+    G_a = pickle.load(f)['G_ema'].to(device) #.cpu()  # torch.nn.Module
 
 cl_text = os.getenv('SHOW_TEXT')
+
+
   
 def gen_image(text):    
     text = text.strip().lower()
@@ -29,13 +34,13 @@ def gen_image(text):
         return gen_image_helper(G_a)
 
 def gen_image_helper(model):
-    z = torch.randn([1, model.z_dim]).cpu()  # latent codes
+    z = torch.randn([1, model.z_dim]).to(device) #.cpu()  # latent codes
     c = None  # class labels (not used in this example)
     img = model(z, c)
     img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
     #um = torch..nn.Upsample(scale_factor=2, mode='bilinear')
     #img=um(img)
-    image=PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+    image=PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB') #.cpu()
     transform = transforms.Resize((image.height * 2, image.width * 2), interpolation=transforms.InterpolationMode.BILINEAR)
     upscaled_image = transform(image)  
     return upscaled_image
@@ -46,7 +51,7 @@ demo = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Prompt here..."),
     outputs=gr.Image(type="pil"),
     title="Text to Image Generator",
-    description="Enter text to generate an image using a custom PyTorch model."
+    description="Enter any text to generate an image of an animal"
 )
 
 if __name__ == "__main__":

network-snapshot.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e4201b5d27d257853336c29943555cbcfbfcda6ec6a25bfbcfb2d17570b1e9c4
+oid sha256:bac188ab58ff51bc7f13ce98b811cde06eed9b428658f9ea99a472ef548b9662
 size 266257697