Pls work

.env

@@ -1 +1,2 @@
-HF_HOME=./models
+HF_HOME=./models
+TRANSFORMERS_CACHE=./cache

Dockerfile.fastapi

@@ -11,4 +11,4 @@ RUN pip install --no-cache-dir uvicorn gunicorn fastapi pytest ruff pytest-async
 
 EXPOSE 80
 
-CMD ["uvicorn", "tld.app:app", "--host", "0.0.0.0", "--port", "80"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80"]

og readme.md

@@ -71,7 +71,7 @@ If you have your own dataset of URLs + captions, the process to train a model on
 ```python
 !wandb login
 import os
-from tld.train import main, DataConfig, ModelConfig
+from train import main, DataConfig, ModelConfig
 from accelerate import notebook_launcher
 
 data_config = DataConfig(latent_path='path/to/image_latents.npy',

old/main.py

@@ -0,0 +1,37 @@
+import os
+from transformers import CLIPProcessor, CLIPModel
+from PIL import Image
+
+# Get the directory of the script
+script_directory = os.path.dirname(os.path.realpath(__file__))
+# Specify the directory where the cache will be stored (same folder as the script)
+cache_directory = os.path.join(script_directory, "cache")
+# Create the cache directory if it doesn't exist
+os.makedirs(cache_directory, exist_ok=True)
+
+# Load the CLIP processor and model
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", cache_dir=cache_directory)
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", cache_dir=cache_directory)
+
+# Text description to generate image
+text = "a cat sitting on a table"
+
+# Tokenize text and get features
+inputs = clip_processor(text, return_tensors="pt", padding=True)
+
+# Generate image from text
+generated_image = clip_model.generate(
+  inputs=inputs.input_ids,
+  attention_mask=inputs.attention_mask,
+  visual_input=None,  # We don't provide image inputvi
+  return_tensors="pt"  # Return PyTorch tensor
+)
+
+# Convert the generated image tensor to a NumPy array
+generated_image_np = generated_image[0].cpu().numpy()
+
+# Save the generated image
+output_image_path = "generated_image.png"
+Image.fromarray(generated_image_np).save(output_image_path)
+
+print("Image generated and saved as:", output_image_path)

requirements.txt

@@ -7,4 +7,5 @@ diffusers
 accelerate
 transformers
 Pillow
+poetry
 git+https://github.com/openai/CLIP.git

start.sh

@@ -1,5 +1,4 @@
 pip install --upgrade pip
 pip install -r requirements.txt
-poetry install --no-root
 
-python main.py
+python tld/gen_img.py

tests/client.js

@@ -0,0 +1,15 @@
+const axios = require("axios");
+
+const apiUrl = `http://de-fsn-4.halex.gg:25287/api`;
+
+const postData = {
+  prompt: "Wassup my homie"
+};
+
+axios.post(apiUrl, postData)
+  .then(response => {
+    console.log("Response from API:", response.data);
+  })
+  .catch(error => {
+    console.error("Error:", error.message);
+  });

tests/test_api.py

@@ -1,7 +1,5 @@
-import os
-
 from fastapi.testclient import TestClient
-from tld.app import app
+from app import app
 import PIL
 from PIL import Image
 from io import BytesIO
@@ -9,23 +7,24 @@ from io import BytesIO
 client = TestClient(app)
 
 def test_read_main():
-    response = client.get("/")
-    assert response.status_code == 200
-    assert response.json() == {"message": "Welcome to Image Generator"}
+  response = client.get("/")
+  assert response.status_code == 200
+  assert response.json() == {"message": "Welcome to Image Generator"}
 
 
 def test_generate_image_unauthorized():
-    response = client.post("/generate-image/", json={})
-    assert response.status_code == 401
-    assert response.json() == {"detail": "Not authenticated"}
+  response = client.post("/generate-image/", json={})
+  assert response.status_code == 401
+  assert response.json() == {"detail": "Not authenticated"}
 
 
 def test_generate_image_authorized():
-    api_token = os.getenv("API_TOKEN")
-    response = client.post(
-        "/generate-image/", json={"prompt": "a cute cat"}, headers={"Authorization": f"Bearer {api_token}"}
-    )
-    assert response.status_code == 200
-
-    image = Image.open(BytesIO(response.content))
-    assert type(image) == PIL.JpegImagePlugin.JpegImageFile
+  response = client.post(
+    "/generate-image/", json={"prompt": "a cute cat"}
+  )
+  assert response.status_code == 200
+
+  image = Image.open(BytesIO(response.content))
+  assert type(image) == PIL.JpegImagePlugin.JpegImageFile
+
+test_generate_image_authorized()

tests/test_diffuser.py

@@ -10,8 +10,8 @@ import torchvision.transforms as transforms
 import torchvision.utils as vutils
 from diffusers import AutoencoderKL
 
-from tld.denoiser import Denoiser
-from tld.diffusion import DiffusionGenerator, DiffusionTransformer, LTDConfig
+from denoiser import Denoiser
+from diffusion import DiffusionGenerator, DiffusionTransformer, LTDConfig
 from PIL.Image import Image
 
 to_pil = transforms.ToPILImage()

tld/app.py

@@ -4,39 +4,37 @@ from typing import Optional
 
 import torch
 import torchvision.transforms as transforms
-from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
 from fastapi.security import OAuth2PasswordBearer
 from pydantic import BaseModel
 
-from tld.diffusion import DiffusionTransformer, LTDConfig
+from diffusion import DiffusionTransformer, LTDConfig
+
+# Get the directory of the script
+script_directory = os.path.dirname(os.path.realpath(__file__))
+# Specify the directory where the cache will be stored (same folder as the script)
+cache_directory = os.path.join(script_directory, "cache")
+home_directory = os.path.join(script_directory, "home")
+# Create the cache directory if it doesn't exist
+os.makedirs(cache_directory, exist_ok=True)
+os.makedirs(home_directory, exist_ok=True)
+
+os.environ["TRANSFORMERS_CACHE"] = cache_directory
+os.environ["HF_HOME"] = home_directory
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 to_pil = transforms.ToPILImage()
-
 ltdconfig = LTDConfig()
-diffusion_transformer = DiffusionTransformer(ltdconfig)
-
+diffusion_transformer = DiffusionTransformer(ltdconfig) #Downloads model here
 app = FastAPI()
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
-
-
-def validate_token(token: str = Depends(oauth2_scheme)):
-    if token != os.getenv("API_TOKEN"):
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid authentication credentials",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-
-
 class ImageRequest(BaseModel):
-    prompt: str
-    class_guidance: Optional[int] = 6
-    seed: Optional[int] = 11
-    num_imgs: Optional[int] = 1
-    img_size: Optional[int] = 32
+  prompt: str
+  class_guidance: Optional[int] = 6
+  seed: Optional[int] = 11
+  num_imgs: Optional[int] = 1
+  img_size: Optional[int] = 32
 
 
 @app.get("/")
@@ -45,23 +43,24 @@ def read_root():
 
 
 @app.post("/generate-image/")
-async def generate_image(request: ImageRequest, token: str = Depends(validate_token)):
-    try:
-        img = diffusion_transformer.generate_image_from_text(
-            prompt=request.prompt,
-            class_guidance=request.class_guidance,
-            seed=request.seed,
-            num_imgs=request.num_imgs,
-            img_size=request.img_size,
-        )
-        # Convert PIL image to byte stream suitable for HTTP response
-        img_byte_arr = io.BytesIO()
-        img.save(img_byte_arr, format="JPEG")
-        img_byte_arr.seek(0)
-
-        return StreamingResponse(img_byte_arr, media_type="image/jpeg")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+async def generate_image(request: ImageRequest):
+  try:
+    img = diffusion_transformer.generate_image_from_text(
+      prompt=request.prompt,
+      class_guidance=request.class_guidance,
+      seed=request.seed,
+      num_imgs=request.num_imgs,
+      img_size=request.img_size,
+    )
+    
+    # Convert PIL image to byte stream suitable for HTTP response
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format="JPEG")
+    img_byte_arr.seek(0)
+
+    return StreamingResponse(img_byte_arr, media_type="image/jpeg")
+  except Exception as e:
+    raise HTTPException(status_code=500, detail=str(e))
 
 
 # build job to test and deploy the API on a docker image (maybe in Azure?)

tld/denoiser.py

@@ -4,7 +4,7 @@ import torch
 from einops.layers.torch import Rearrange
 from torch import nn
 
-from tld.transformer_blocks import DecoderBlock, MLPSepConv, SinusoidalEmbedding
+from transformer_blocks import DecoderBlock, MLPSepConv, SinusoidalEmbedding
 
 
 class DenoiserTransBlock(nn.Module):

tld/diffusion.py

@@ -10,7 +10,7 @@ from diffusers import AutoencoderKL
 from torch import Tensor
 from tqdm import tqdm
 
-from tld.denoiser import Denoiser
+from denoiser import Denoiser
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 to_pil = transforms.ToPILImage()
@@ -124,21 +124,21 @@ class DiffusionGenerator:
 
 @dataclass
 class LTDConfig:
-    vae_scale_factor: float = 8
-    img_size: int = 32
-    model_dtype: torch.dtype = torch.float32
-    file_url: str = None  # = "https://huggingface.co/apapiu/small_ldt/resolve/main/state_dict_378000.pth"
-    local_filename: str = "state_dict_378000.pth"
-    vae_name: str = "madebyollin/sdxl-vae-fp16-fix"
-    clip_model_name: str = "ViT-L/14"
-    denoiser: Denoiser = Denoiser(
-        image_size=32,
-        noise_embed_dims=256,
-        patch_size=2,
-        embed_dim=256,
-        dropout=0,
-        n_layers=4,
-    )
+  vae_scale_factor: float = 8
+  img_size: int = 32
+  model_dtype: torch.dtype = torch.float32
+  file_url: str = None  # = "https://huggingface.co/apapiu/small_ldt/resolve/main/state_dict_378000.pth"
+  local_filename: str = "state_dict_378000.pth"
+  vae_name: str = "ByteDance/SDXL-Lightning"
+  clip_model_name: str = "ViT-L/14"
+  denoiser: Denoiser = Denoiser(
+    image_size=32,
+    noise_embed_dims=256,
+    patch_size=2,
+    embed_dim=256,
+    dropout=0,
+    n_layers=4,
+  )
 
 
 def download_file(url, filename):
@@ -151,48 +151,48 @@ def download_file(url, filename):
 
 @torch.no_grad()
 def encode_text(label, model):
-    text_tokens = clip.tokenize(label, truncate=True).to(device)
-    text_encoding = model.encode_text(text_tokens)
-    return text_encoding.cpu()
+  text_tokens = clip.tokenize(label, truncate=True).to(device)
+  text_encoding = model.encode_text(text_tokens)
+  return text_encoding.cpu()
 
 
 class DiffusionTransformer:
-    def __init__(self, config: LTDConfig):
-        denoiser = config.denoiser.to(config.model_dtype)
-
-        if config.file_url is not None:
-            print(f"Downloading model from {config.file_url}")
-            download_file(config.file_url, config.local_filename)
-            state_dict = torch.load(config.local_filename, map_location=torch.device("cpu"))
-            denoiser.load_state_dict(state_dict)
-
-        denoiser = denoiser.to(device)
-
-        vae = AutoencoderKL.from_pretrained(config.vae_name, torch_dtype=config.model_dtype).to(device)
-
-        self.clip_model, preprocess = clip.load(config.clip_model_name)
-        self.clip_model = self.clip_model.to(device)
-
-        self.diffuser = DiffusionGenerator(denoiser, vae, device, config.model_dtype)
-
-    def generate_image_from_text(
-        self, prompt: str, class_guidance=6, seed=11, num_imgs=1, img_size=32, n_iter=15
-    ):
-        nrow = int(np.sqrt(num_imgs))
-
-        cur_prompts = [prompt] * num_imgs
-        labels = encode_text(cur_prompts, self.clip_model)
-        out, out_latent = self.diffuser.generate(
-            labels=labels,
-            num_imgs=num_imgs,
-            class_guidance=class_guidance,
-            seed=seed,
-            n_iter=n_iter,
-            exponent=1,
-            scale_factor=8,
-            sharp_f=0,
-            bright_f=0,
-        )
+  def __init__(self, config: LTDConfig):
+    denoiser = config.denoiser.to(config.model_dtype)
+
+    if config.file_url is not None:
+      print(f"Downloading model from {config.file_url}")
+      download_file(config.file_url, config.local_filename)
+      state_dict = torch.load(config.local_filename, map_location=torch.device("cpu"))
+      denoiser.load_state_dict(state_dict)
+
+    denoiser = denoiser.to(device)
+
+    vae = AutoencoderKL.from_pretrained(config.vae_name, torch_dtype=config.model_dtype).to(device)
+
+    self.clip_model, preprocess = clip.load(config.clip_model_name)
+    self.clip_model = self.clip_model.to(device)
+
+    self.diffuser = DiffusionGenerator(denoiser, vae, device, config.model_dtype)
+
+  def generate_image_from_text(
+    self, prompt: str, class_guidance=6, seed=11, num_imgs=1, img_size=32, n_iter=15
+  ):
+    nrow = int(np.sqrt(num_imgs))
+
+    cur_prompts = [prompt] * num_imgs
+    labels = encode_text(cur_prompts, self.clip_model)
+    out, out_latent = self.diffuser.generate(
+      labels=labels,
+      num_imgs=num_imgs,
+      class_guidance=class_guidance,
+      seed=seed,
+      n_iter=n_iter,
+      exponent=1,
+      scale_factor=8,
+      sharp_f=0,
+      bright_f=0,
+    )
 
-        out = to_pil((vutils.make_grid((out + 1) / 2, nrow=nrow, padding=4)).float().clip(0, 1))
-        return out
+    out = to_pil((vutils.make_grid((out + 1) / 2, nrow=nrow, padding=4)).float().clip(0, 1))
+    return out

tld/gen_img.py

@@ -0,0 +1,44 @@
+import io
+import asyncio
+import os
+from PIL import Image
+import numpy as np
+
+import torch
+import torchvision.transforms as transforms
+from torchvision import utils as vutils
+
+from diffusion import DiffusionTransformer, LTDConfig
+
+# Get the directory of the script
+script_directory = os.path.dirname(os.path.realpath(__file__))
+# Specify the directory where the cache will be stored (same folder as the script)
+cache_directory = os.path.join(script_directory, "cache")
+home_directory = os.path.join(script_directory, "home")
+# Create the cache directory if it doesn't exist
+os.makedirs(cache_directory, exist_ok=True)
+os.makedirs(home_directory, exist_ok=True)
+
+os.environ["TRANSFORMERS_CACHE"] = cache_directory
+os.environ["HF_HOME"] = home_directory
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+to_pil = transforms.ToPILImage()
+ltdconfig = LTDConfig()
+diffusion_transformer = DiffusionTransformer(ltdconfig) #Downloads model here
+
+async def generate_image(prompt):
+  try:
+    img = diffusion_transformer.generate_image_from_text(
+      prompt=prompt,
+      class_guidance=6,
+      seed=11,
+      num_imgs=1,
+      img_size=32,
+    )
+    
+    img.save("generated_img.png")
+  except Exception as e:
+    print(e)
+    
+asyncio.run(generate_image("a cute cat"))

tld/train.py

@@ -15,8 +15,8 @@ from torch import Tensor, nn
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm import tqdm
 
-from tld.denoiser import Denoiser
-from tld.diffusion import DiffusionGenerator
+from denoiser import Denoiser
+from diffusion import DiffusionGenerator
 
 
 def eval_gen(diffuser: DiffusionGenerator, labels: Tensor) -> Image: