Attempt at fixing model

app.py

@@ -7,11 +7,11 @@ from transformers import BertTokenizer, BertModel
 import numpy as np
 import os
 import time
+from typing import Optional, Union
 
 LATENT_DIM = 128
 HIDDEN_DIM = 256
 
-
 # Text encoder
 class TextEncoder(nn.Module):
     def __init__(self, hidden_size, output_size):
@@ -23,7 +23,7 @@ class TextEncoder(nn.Module):
         outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
         return self.fc(outputs.last_hidden_state[:, 0, :])
 
-# CVAE model
+# CVAE model (unchanged)
 class CVAE(nn.Module):
     def __init__(self, text_encoder):
         super(CVAE, self).__init__()
@@ -81,14 +81,20 @@ class CVAE(nn.Module):
 # Initialize the BERT tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
-def clean_image(image, threshold=0.75):
+def clean_image(image: Image.Image, threshold: float = 0.75) -> Image.Image:
     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
     return Image.fromarray(np_image)
 
-def generate_image(model, text_prompt, device, input_image=None, img_control=0.5):
+def generate_image(
+    model: CVAE,
+    text_prompt: str,
+    device: torch.device,
+    input_image: Optional[Image.Image] = None,
+    img_control: float = 0.5
+) -> Image.Image:
     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)
@@ -110,31 +116,52 @@ def generate_image(model, text_prompt, device, input_image=None, img_control=0.5
 
     return generated_image
 
-def load_model(model_path, device):
-    text_encoder = TextEncoder(hidden_size=HIDDEN_DIM, output_size=HIDDEN_DIM)
-    model = CVAE(text_encoder).to(device)
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.eval()
-    return model
-
-def generate_image_gradio(prompt, model_path, clean_image_flag, size, input_image=None, img_control=0.5):
+# Model loading with caching
+_model_cache = {}
+def load_model(model_path: str, device: torch.device) -> CVAE:
+    if model_path not in _model_cache:
+        text_encoder = TextEncoder(hidden_size=HIDDEN_DIM, output_size=HIDDEN_DIM)
+        model = CVAE(text_encoder).to(device)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.eval()
+        _model_cache[model_path] = model
+    return _model_cache[model_path]
+
+def generate_image_gradio(
+    prompt: str,
+    model_path: str,
+    clean_image_flag: bool,
+    size: int,
+    input_image: Optional[Image.Image] = None,
+    img_control: float = 0.5
+) -> tuple[Image.Image, str]:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = load_model(model_path, device)
+    
+    try:
+        model = load_model(model_path, device)
+    except Exception as e:
+        raise gr.Error(f"Failed to load model: {str(e)}")
 
     start_time = time.time()
-    generated_image = generate_image(model, prompt, device, input_image, img_control)
+    try:
+        generated_image = generate_image(model, prompt, device, input_image, img_control)
+    except Exception as e:
+        raise gr.Error(f"Failed to generate image: {str(e)}")
+    
     end_time = time.time()
     generation_time = end_time - start_time
 
     if clean_image_flag:
         generated_image = clean_image(generated_image)
 
-    generated_image = generated_image.resize((size, size), resample=Image.NEAREST)
+    try:
+        generated_image = generated_image.resize((size, size), resample=Image.NEAREST)
+    except Exception as e:
+        raise gr.Error(f"Failed to resize image: {str(e)}")
 
     return generated_image, f"Generation time: {generation_time:.4f} seconds"
 
-# Gradio interface
-def gradio_interface():
+def gradio_interface() -> gr.Blocks:
     with gr.Blocks() as demo:
         gr.Markdown("# Image Generator from Text Prompt")
         
@@ -152,14 +179,23 @@ def gradio_interface():
                 output_image = gr.Image(label="Generated Image")
                 generation_time = gr.Textbox(label="Generation Time")
 
+        # Use gr.Error for error handling
         generate_button.click(
-            generate_image_gradio,
+            fn=generate_image_gradio,
             inputs=[prompt, model_path, clean_image_flag, size, input_image, img_control],
-            outputs=[output_image, generation_time]
+            outputs=[output_image, generation_time],
+            api_name="generate"  # Explicit API endpoint name
         )
 
     return demo
 
 if __name__ == "__main__":
     demo = gradio_interface()
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        # Configure CORS if needed
+        # allowed_paths=["/custom/path"],
+        # cors_allowed_origins=["*"]
+    )