Update app.py

app.py

@@ -3,22 +3,17 @@ import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
 from transformers import BertTokenizer, BertModel
-import argparse
+import gradio as gr
 import numpy as np
 import os
-import time  # Import the time module
+import time
 
-# 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
@@ -26,19 +21,15 @@ def clean_image(image, threshold=0.75):
     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)
     
@@ -47,7 +38,6 @@ def generate_image(model, text_prompt, device, input_image=None, img_control=0.5
         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)
@@ -55,86 +45,62 @@ def generate_image(model, text_prompt, device, input_image=None, img_control=0.5
     
     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]
-
+def process(prompt, model_path, clean, size, input_image, img_control, output_dir):
     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]
+    input_image = Image.open(input_image).convert("RGBA") if input_image else None
+    
+    # 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()
 
-        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
+    start_time = time.time()
+    
+    # Generate image from prompt
+    generated_image = generate_image(model, prompt, device, input_image, img_control)
+    
+    end_time = time.time()
+    generation_time = end_time - start_time
+    
+    # Clean up the image if the flag is set
+    if clean:
+        generated_image = clean_image(generated_image)
+    
+    # Resize the generated image
+    generated_image = generated_image.resize((size, size), resample=Image.NEAREST)
+    
+    # Save the generated image to the specified directory
+    model_name = os.path.splitext(os.path.basename(model_path))[0]
+    output_file = os.path.join(output_dir, f"{model_name}_{prompt}.png")
+    os.makedirs(output_dir, exist_ok=True)
+    generated_image.save(output_file)
+    
+    print(f"Generated image saved as {output_file}")
+    print(f"Generation time: {generation_time:.10f} seconds")
+    
+    return generated_image
 
-            # 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
+# Gradio Interface
+interface = gr.Interface(
+    fn=process,
+    inputs=[
+        gr.Textbox(label="Text Prompt"),
+        gr.File(label="Model Path (.pth file)", file_types=['.pth']),
+        gr.Checkbox(label="Clean Image (Remove Low Opacity Pixels)", default=False),
+        gr.Slider(label="Image Size", minimum=16, maximum=512, step=16, default=16),
+        gr.File(label="Input Image (Optional)", file_types=["image"]),
+        gr.Slider(label="Image Control (0-1)", minimum=0.0, maximum=1.0, step=0.01, default=0.5),
+        gr.Textbox(label="Output Directory", value="generated_images")
+    ],
+    outputs=gr.Image(label="Generated Image"),
+    title="Text-to-Image Generator",
+    description="Generate an image from a text prompt using a trained CVAE model."
+)
 
 if __name__ == "__main__":
-    main()
+    interface.launch()