Update app.py

app.py

@@ -1,41 +1,60 @@
-import io
 import requests
+import json
 from PIL import Image
-import torch
-from torchvision import transforms
-from torchvision.utils import save_image
-from torchvision.transforms.functional import resize
-from argparse import Namespace
 
-# Load pre-trained StyleGAN2 model from Hugging Face
-model = torch.hub.load('huggingface/pytorch-pretrained-StyleGAN2', 'StyleGAN2')
+# Replace YOUR_API_KEY with your OpenAI API key
+API_KEY = 'YOUR_API_KEY'
+ENDPOINT = 'https://api.openai.com/v1/images/generations'
 
-# Set device to GPU if available
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-
-# Set seed for reproducibility
-torch.manual_seed(0)
-
-# Define function to generate image from text prompt
 def generate_image(prompt):
-    # Convert text prompt to embeddings using StyleGAN2's tokenizer
-    embeddings = model.gen.encode(prompt)
-
-    # Generate image from embeddings
-    with torch.no_grad():
-        generated_image = model.gen.decode(embeddings, input_is_latent=True)
-
-    # Resize and save image
-    generated_image = generated_image.clamp(0, 1)
-    generated_image = resize(generated_image, (256, 256), interpolation=Image.LANCZOS)
-    image_tensor = transforms.ToTensor()(generated_image)
-    image_tensor = image_tensor.unsqueeze(0)
-    save_image(image_tensor, 'output.png')
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {API_KEY}'
+    }
+
+    data = {
+        'model': 'image-alpha-001',
+        'prompt': prompt,
+        'num_images': 1,
+        'size': '512x512',
+        'response_format': 'url'
+    }
+
+    response = requests.post(url=ENDPOINT, headers=headers, data=json.dumps(data))
+    if response.status_code == 200:
+        result_url = response.json()['data'][0]['url']
+        image = Image.open(requests.get(result_url, stream=True).raw)
+        return image
+    else:
+        return None
 
 # Example usage
-generate_image("A peaceful countryside with a lake and mountains in the distance")
-
-# Load generated image and display
-image = Image.open('output.png')
-image.show()
+prompt = "a cat sitting on a windowsill looking outside"
+image = generate_image(prompt)
+if image:
+    image.show()
+else:
+    print("Error generating image")
+# A red sports car driving on a winding road
+prompt = "a red sports car driving on a winding road"
+image = generate_image(prompt)
+if image:
+    image.show()
+else:
+    print("Error generating image")
+
+# A house in a forest with a river nearby
+prompt = "a house in a forest with a river nearby"
+image = generate_image(prompt)
+if image:
+    image.show()
+else:
+    print("Error generating image")
+
+# A slice of pepperoni pizza on a plate
+prompt = "a slice of pepperoni pizza on a plate"
+image = generate_image(prompt)
+if image:
+    image.show()
+else:
+    print("Error generating image")