Update README.md

README.md

@@ -1,83 +1,102 @@
----
-license: mit
----
+# ArtifyAI v1.0: Text-to-Image Generation
 
-# ArtifyAI - COCO Dataset Downloader and Organizer
+ArtifyAI v1.0 is a project designed to generate images from text using the T5 model for text summarization and the Stable Diffusion model for image generation. This version also provides functionality to save model weights and load them for future use, making it easier to resume your work.
 
-Welcome to ArtifyAI, a project designed to download images from the COCO dataset and organize them into your Google Drive for further use in machine learning and computer vision projects.
+## Overview
 
-## Introduction
+ArtifyAI v1.0 uses T5 for text processing, Stable Diffusion for image generation, and includes functionality for saving and loading the model's weights. You can easily store model weights in Google Drive and use them later without re-downloading everything.
 
-This project automates the process of downloading a subset of images from the COCO dataset and storing them directly in Google Drive. It's aimed at researchers and developers who need a quick way to set up their dataset for training image-based models. The notebook is structured to be easy to follow, even for non-technical users.
+## Features
 
-## Requirements
+- **Text Processing**: Utilizes T5 for text summarization and generation.
+- **Image Generation**: Uses Stable Diffusion for high-quality image generation based on textual descriptions.
+- **Model Saving**: Allows you to save and reload model weights (e.g., UNet) to/from Google Drive.
 
-To use this notebook, you need:
+## Installation
 
-- Google Colab (optional, but recommended for ease of use)
-- Google Drive (to store the dataset)
-- Basic knowledge of Python (optional)
-- Internet connection (for downloading images)
-- Libraries:
-  - `pycocotools`
-  - `requests`
-  - `tqdm`
-  - `matplotlib`
-  - `numpy`
-  - `shutil`
+### Prerequisites
 
-## Installation
+Ensure that you have the following to run this project:
 
-1. **Clone the repository (optional)**:
-   ```bash
-   git clone https://github.com/your-repo/ArtifyAI.git
-   cd ArtifyAI
-   ```
+1. Python 3.7+
+2. CUDA-compatible GPU (for faster performance)
+3. [Hugging Face Transformers](https://huggingface.co/transformers/)
+4. [Diffusers](https://huggingface.co/docs/diffusers/index)
+5. [PyTorch](https://pytorch.org/) with CUDA support
 
-2. **Open the notebook**:
-   - You can directly upload the notebook (`ArtifyAI_v1_0.ipynb`) to your Google Colab environment or use it locally in your Python environment.
+### Step-by-Step Setup
 
-3. **Install dependencies**:
-   If you're running this locally, make sure to install the required libraries:
-   ```bash
-   pip install pycocotools tqdm matplotlib numpy requests
-   ```
+1. **Clone the Repository**:
+    ```bash
+    git clone https://github.com/your-username/ArtifyAI.git
+    cd ArtifyAI
+    ```
 
-## Running the Project
+2. **Install Dependencies**:
+    ```bash
+    pip install torch transformers diffusers
+    ```
 
-Once all dependencies are installed, follow these steps:
+3. **Download the Pretrained Models**:
+    Use the following code to load the models:
+    ```python
+    from transformers import T5Tokenizer, T5ForConditionalGeneration
+    from diffusers import StableDiffusionPipeline
+    import torch
 
-1. **Open Google Colab**:
-   - If you're using Google Colab, upload the notebook to your Colab environment.
-   
-2. **Mount Google Drive**:
-   - The notebook will prompt you to mount your Google Drive. Ensure your account is connected so that the images are downloaded and saved to a designated folder in your drive.
+    # Load models
+    t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
+    t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
+    ArtifyAI_model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
+
+    # Set model to GPU if available
+    ArtifyAI_model = ArtifyAI_model.to("cuda" if torch.cuda.is_available() else "cpu")
+    ```
 
-3. **Run the Code**:
-   - Execute the cells sequentially to start downloading the COCO dataset images. The notebook uses `pycocotools` to fetch image URLs from the COCO API and downloads them using `requests`. Progress is tracked with `tqdm`.
+4. **Mount Google Drive for Model Saving**:
+    To save model weights in Google Drive, mount your drive and specify the save path:
+    ```python
+    from google.colab import drive
+    drive.mount('/content/drive')
 
-4. **Transfer Images to Google Drive**:
-   - After downloading, the images will automatically be moved to a folder in your Google Drive (`MyDrive/coco_dataset`).
+    save_path = '/content/drive/My Drive/models/unet_weights.pth'
+    torch.save(ArtifyAI_model.unet.state_dict(), save_path)
+    ```
 
-## Using Google Colab
+5. **Load Model Weights**:
+    Reload the model's weights later using:
+    ```python
+    ArtifyAI_model.unet.load_state_dict(torch.load('/content/drive/My Drive/models/unet_weights.pth'))
+    ```
 
-For non-technical users, we recommend using [Google Colab](https://colab.research.google.com/). It provides a cloud-based environment where you can run the notebook without installing Python or any dependencies on your local machine.
+## Usage
 
-1. **Upload the Notebook**:
-   - Simply drag and drop the `.ipynb` file into Colab.
+1. **Run the Jupyter Notebook**: Open `ArtifyAI_v1_0.ipynb` in Jupyter to run the code and explore the pipeline.
+2. **Save/Load Model Weights**: Use the provided code to store and retrieve model weights from Google Drive.
+3. **Generate Custom Images**: Modify the text input to generate unique images.
 
-2. **Mount Google Drive**:
-   - The notebook includes a step to mount Google Drive for file storage. Follow the on-screen instructions to authorize access.
+## Example
 
-3. **Run the Notebook**:
-   - Execute each cell in the notebook by clicking the play button next to each cell. Ensure all code cells are run in order.
+Here's how you can save and reload model weights:
+```python
+# Save the UNet model's weights
+torch.save(ArtifyAI_model.unet.state_dict(), '/content/drive/My Drive/models/unet_weights.pth')
 
-## Moving Files to Google Drive
+# Load the saved weights
+ArtifyAI_model.unet.load_state_dict(torch.load('/content/drive/My Drive/models/unet_weights.pth'))
 
-After downloading the images, the notebook will move them to a specified folder in your Google Drive, making it easy for you to access them later.
+## For Non-Technical Users
+
+If you're not familiar with coding or AI, follow these steps to use ArtifyAI:
+
+1. **Install Python**: Download and install Python 3.7+ from the [official Python website](https://www.python.org/downloads/).
+   
+2. **Install Dependencies**: After installing Python, open a terminal and run:
+    ```bash
+    pip install torch transformers diffusers
+    ```
 
-- The default folder is `/content/drive/MyDrive/coco_dataset`, but you can modify the path if needed.
+3. **Run the Project**: Use the code snippets provided in the notebook or above to generate images from text.
 
-## License
+4. **Save Your Work**: If you are using Google Colab, remember to save your models to Google Drive to keep your work.
 
-This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.