Fixed app.py

app.py

@@ -10,26 +10,47 @@ import pickle
 import torchvision.transforms as transforms
 from PIL import Image
 
-network_pkl = 'network-snapshot-003200.pkl'
-with open(network_pkl, 'rb') as f:
-    G = pickle.load(f)['G_ema'].cpu()  # torch.nn.Module
+network_pkl_a = 'stylegan3-r-afhqv2-512x512.pkl'
+network_pkl_d = 'network-snapshot.pkl'
+with open(network_pkl_d, 'rb') as f:
+    G = pickle.load(f)['G_ema']  # torch.nn.Module
 
-def gen_image(text_input):    
-    z = torch.randn([1, G.z_dim]).cpu()  # latent codes
+with open(network_pkl_a, 'rb') as f:
+    G_a = pickle.load(f)['G_ema']  # torch.nn.Module
+    
+def gen_image(text):
+    if text=='show me':
+        return gen_image_d()
+    else:
+        return gen_image_a()
+        
+def gen_image_a():    
+    z = torch.randn([1, G_a.z_dim])  # latent codes
+    c = None  # class labels (not used in this example)
+    img = G_a(z, c)
+    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    image=PIL.Image.fromarray(img[0].numpy(), 'RGB')
+    transform = transforms.Resize((image.height * 2, image.width * 2), interpolation=transforms.InterpolationMode.BILINEAR)
+    upscaled_image = transform(image)
+    return upscaled_image
+
+def gen_image_d():    
+    z = torch.randn([1, G.z_dim])  # latent codes
     c = None  # class labels (not used in this example)
     img = G(z, c)
     img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-    image=PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+    image=PIL.Image.fromarray(img[0].numpy(), 'RGB')
     transform = transforms.Resize((image.height * 2, image.width * 2), interpolation=transforms.InterpolationMode.BILINEAR)
     upscaled_image = transform(image)
     return upscaled_image
     
 demo = gr.Interface(
     fn=gen_image,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
+    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."
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

stylegan3-r-afhqv2-512x512.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77b9cd6bbaf0a2dfeb4372c9f86b7280a98f293aa1b9d35465864895596628c8
+size 249525556