Create app.py

app.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+from transformers import BertTokenizer, BertModel
+import argparse
+import numpy as np
+import os
+import time  # Import the time module
+
+# Import the model architecture from train.py
+from train import CVAE, TextEncoder, LATENT_DIM, HIDDEN_DIM
+
+# Initialize the BERT tokenizer
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
+def clean_image(image, threshold=0.75):
+    """
+    Clean up the image by setting pixels with opacity <= threshold to 0% opacity
+    and pixels above the threshold to 100% visibility.
+    """
+    np_image = np.array(image)
+    alpha_channel = np_image[:, :, 3]
+    alpha_channel[alpha_channel <= int(threshold * 255)] = 0
+    alpha_channel[alpha_channel > int(threshold * 255)] = 255  # Set to 100% visibility
+    return Image.fromarray(np_image)
+
+def generate_image(model, text_prompt, device, input_image=None, img_control=0.5):
+    # Encode text prompt using BERT tokenizer
+    encoded_input = tokenizer(text_prompt, padding=True, truncation=True, return_tensors="pt")
+    input_ids = encoded_input['input_ids'].to(device)
+    attention_mask = encoded_input['attention_mask'].to(device)
+    
+    # Generate text encoding
+    with torch.no_grad():
+        text_encoding = model.text_encoder(input_ids, attention_mask)
+    
+    # Sample from the latent space
+    z = torch.randn(1, LATENT_DIM).to(device)
+    
+    # Generate image
+    with torch.no_grad():
+        generated_image = model.decode(z, text_encoding)
+    
+    if input_image is not None:
+        input_image = input_image.convert("RGBA").resize((16, 16), resample=Image.NEAREST)
+        input_image = transforms.ToTensor()(input_image).unsqueeze(0).to(device)
+        generated_image = img_control * input_image + (1 - img_control) * generated_image
+    
+    # Convert the generated tensor to a PIL Image
+    generated_image = generated_image.squeeze(0).cpu()
+    generated_image = (generated_image + 1) / 2  # Rescale from [-1, 1] to [0, 1]
+    generated_image = generated_image.clamp(0, 1)
+    generated_image = transforms.ToPILImage()(generated_image)
+    
+    return generated_image
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate an image from a text prompt using the trained CVAE model(s).")
+    parser.add_argument("--prompt", type=str, help="Text prompt for image generation")
+    parser.add_argument("--prompt_file", type=str, help="File containing prompts, one per line")
+    parser.add_argument("--output", type=str, default="generated_images", help="Output directory or file for generated images")
+    parser.add_argument("--model_paths", type=str, nargs='*', help="Paths to the trained model(s)")
+    parser.add_argument("--model_path", type=str, help="Path to a single trained model")
+    parser.add_argument("--clean", action="store_true", help="Clean up the image by removing low opacity pixels")
+    parser.add_argument("--size", type=int, default=16, help="Size of the generated image")
+    parser.add_argument("--input_image", type=str, help="Path to the input image for img2img generation")
+    parser.add_argument("--img_control", type=float, default=0.5, help="Control how much the input image influences the output (0 to 1)")
+    args = parser.parse_args()
+
+    if not args.prompt and not args.prompt_file:
+        parser.error("Either --prompt or --prompt_file must be provided")
+    
+    if args.model_paths and args.model_path:
+        parser.error("Specify either --model_paths or --model_path, not both")
+    
+    model_paths = args.model_paths if args.model_paths else [args.model_path]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Check if --output is a file or directory
+    is_folder_output = os.path.isdir(args.output)
+
+    if is_folder_output:
+        # Ensure output directory exists if it's not a file
+        os.makedirs(args.output, exist_ok=True)
+
+    # Load input image if provided
+    input_image = None
+    if args.input_image:
+        input_image = Image.open(args.input_image).convert("RGBA")
+
+    # Process single prompt or batch of prompts
+    if args.prompt:
+        prompts = [args.prompt]
+    else:
+        with open(args.prompt_file, 'r') as f:
+            prompts = [line.strip() for line in f if line.strip()]
+
+    for model_path in model_paths:
+        # Initialize model
+        text_encoder = TextEncoder(hidden_size=HIDDEN_DIM, output_size=HIDDEN_DIM)
+        model = CVAE(text_encoder).to(device)
+        
+        # Load the trained model
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.eval()
+
+        model_name = os.path.splitext(os.path.basename(model_path))[0]
+
+        for i, prompt in enumerate(prompts):
+            start_time = time.time()  # Start timing the generation
+            
+            # Generate image from prompt
+            generated_image = generate_image(model, prompt, device, input_image, args.img_control)
+            
+            # End timing the generation
+            end_time = time.time()
+            generation_time = end_time - start_time  # Calculate the generation time
+
+            # Clean up the image if the flag is set
+            if args.clean:
+                generated_image = clean_image(generated_image)
+            
+            # Resize the generated image
+            generated_image = generated_image.resize((args.size, args.size), resample=Image.NEAREST)
+            
+            if not is_folder_output:
+                # Save the generated image to the specified file
+                output_file = args.output
+            else:
+                # Save the generated image to the output directory
+                output_file = os.path.join(args.output, f"{model_name}_{prompt}_{i:03d}.png")
+            
+            generated_image.save(output_file)
+            print(f"Generated image for prompt '{prompt}' using model '{model_name}' saved as {output_file}")
+            print(f"Generation time: {generation_time:.10f} seconds")  # Print the generation time
+
+if __name__ == "__main__":
+    main()