Update app.py

app.py

@@ -1,122 +1,66 @@
+import gradio as gr
 import torch
-from torch import nn, optim
-from torch.utils.data import DataLoader, Dataset
-from torchvision import transforms, datasets, models
-from PIL import Image
-import json
+from diffusers import StableDiffusion3Pipeline
 import os
-import gradio as gr
-
-# Paths
-image_folder = "Images/"
-metadata_file = "descriptions.json"
-
-# Define the function to load metadata
-def load_metadata(metadata_file):
-    with open(metadata_file, 'r') as f:
-        metadata = json.load(f)
-    return metadata
-
-# Custom Dataset Class
-class ImageDescriptionDataset(Dataset):
-    def __init__(self, image_folder, metadata):
-        self.image_folder = image_folder
-        self.metadata = metadata
-        self.image_names = list(metadata.keys())  # List of image filenames
-        self.transform = transforms.Compose([
-            transforms.Resize((512, 512)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
-
-    def __len__(self):
-        return len(self.image_names)
-
-    def __getitem__(self, idx):
-        image_name = self.image_names[idx]
-        image_path = os.path.join(self.image_folder, image_name)
-        image = Image.open(image_path).convert("RGB")
-        description = self.metadata[image_name]
-        image = self.transform(image)
-        return image, description
-
-# LoRA Layer Implementation
-class LoRALayer(nn.Module):
-    def __init__(self, original_layer, rank=4):
-        super(LoRALayer, self).__init__()
-        self.original_layer = original_layer
-        self.rank = rank
-        self.lora_up = nn.Linear(original_layer.in_features, rank, bias=False)
-        self.lora_down = nn.Linear(rank, original_layer.out_features, bias=False)
-
-    def forward(self, x):
-        return self.original_layer(x) + self.lora_down(self.lora_up(x))
-
-# LoRA Model Class
-class LoRAModel(nn.Module):
-    def __init__(self):
-        super(LoRAModel, self).__init__()
-        self.backbone = models.resnet18(pretrained=True)  # Base model
-        self.backbone.fc = LoRALayer(self.backbone.fc)  # Replace the final layer with LoRA
-
-    def forward(self, x):
-        return self.backbone(x)
-
-# Training Function
-def train_lora(image_folder, metadata):
-    print("Starting LoRA training process...")
-
-    # Create dataset and dataloader
-    dataset = ImageDescriptionDataset(image_folder, metadata)
-    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
-
-    # Initialize model, loss function, and optimizer
-    model = LoRAModel()
-    criterion = nn.CrossEntropyLoss()  # Update this if your task changes
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
-
-    # Training loop
-    num_epochs = 5  # Adjust the number of epochs based on your needs
-    for epoch in range(num_epochs):
-        print(f"Epoch {epoch + 1}/{num_epochs}")
-        for batch_idx, (images, descriptions) in enumerate(dataloader):
-            # Convert descriptions to a numerical format (if applicable)
-            labels = torch.randint(0, 100, (images.size(0),))  # Placeholder labels
-
-            # Forward pass
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            
-            # Backward pass
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            
-            if batch_idx % 10 == 0:  # Log every 10 batches
-                print(f"Batch {batch_idx}, Loss: {loss.item()}")
-
-    # Save the trained model
-    model_path = "lora_model.pth"
-    torch.save(model.state_dict(), model_path)
-    print(f"Model saved as {model_path}")
-
-    print("Training completed.")
-    return model_path  # Return the path of the saved model
-
-# Gradio App
-def start_training_gradio():
-    print("Loading metadata and preparing dataset...")
-    metadata = load_metadata(metadata_file)
-    model_path = train_lora(image_folder, metadata)
-    return model_path  # This will return the model file path for download
-
-# Gradio interface
-demo = gr.Interface(
-    fn=start_training_gradio,
-    inputs=None,
-    outputs=gr.File(),
-    title="Train LoRA Model",
-    description="Fine-tune a model using LoRA for consistent image generation."
-)
-
-demo.launch()
+import spaces
+
+# Use the token saved in secrets
+hf_token = os.getenv("HF_TOKEN")
+
+# Specify the pre-trained model ID
+model_id = "stabilityai/stable-diffusion-3.5-large"
+
+# Global variable for the pipeline (only initialized once)
+pipeline = None
+
+# Function for initializing and caching the pipeline
+def initialize_pipeline():
+    global pipeline
+    if pipeline is None:
+        try:
+            # Load the pipeline with mixed precision (FP16)
+            pipeline = StableDiffusion3Pipeline.from_pretrained(
+                model_id,
+                use_auth_token=hf_token,
+                torch_dtype=torch.float16,  # Use FP16 for mixed precision
+            )
+            # Enable model offloading and attention slicing for memory efficiency
+            pipeline.enable_model_cpu_offload()
+            pipeline.enable_attention_slicing()
+            print("Pipeline initialized and cached.")
+        except Exception as e:
+            # Error handling for model loading issues
+            print(f"Error loading the model: {e}")
+            raise RuntimeError("Failed to initialize the model pipeline.")
+    return pipeline
+
+# Function for image generation, decorated to use GPU
+@spaces.GPU(duration=65)
+def generate_image(prompt):
+    pipe = initialize_pipeline()  # Initialize the pipeline (only once)
+    # Generate the image using the pipeline
+    try:
+        image = pipe(prompt).images[0]
+    except Exception as e:
+        # Catch errors during image generation (e.g., GPU/Memory errors)
+        print(f"Error during image generation: {e}")
+        raise RuntimeError("Image generation failed.")
+    return image
+
+# Set up Gradio interface with a simple input for text and output for image
+interface = gr.Interface(fn=generate_image, inputs="text", outputs="image")
+
+# Launch the interface
+interface.launch()
+
+# Optimize device and dtype handling for CUDA or CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+# Additional model validation (this is optional, more for debugging)
+pipe = initialize_pipeline()  # Ensure the model is initialized and cached
+if not pipe or not hasattr(pipe, 'transformer'):
+    raise ValueError("Failed to load the model or the transformer component is missing.")
+
+# Move the pipeline to the correct device (CUDA or CPU)
+pipe = pipe.to(device)