Update app.py

app.py

@@ -3,18 +3,24 @@ import gradio as gr
 import json
 import logging
 import torch
-from os import path
 from PIL import Image
+from os import path
+from torchvision import transforms
+from dataclasses import dataclass
+import math
+from typing import Callable
 import spaces
-from diffusers import DiffusionPipeline, AutoPipelineForText2Image
-from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler, SD3Transformer2DModel # pip install diffusers>=0.31.0
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL, AutoPipelineForText2Image
+from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler # pip install diffusers>=0.31.0
 from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
+from diffusers.models.transformers import SD3Transformer2DModel
 import copy
 import random
 import time
+import safetensors.torch
+from tqdm import tqdm
 from huggingface_hub import HfFileSystem, ModelCard
 from huggingface_hub import login, hf_hub_download
-import safetensors.torch
 from safetensors.torch import load_file
 hf_token = os.environ.get("HF_TOKEN")
 login(token=hf_token)
@@ -24,7 +30,7 @@ os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 
-torch.set_float32_matmul_precision("medium")
+#torch.set_float32_matmul_precision("medium")
 
 #torch._inductor.config.conv_1x1_as_mm = True
 #torch._inductor.config.coordinate_descent_tuning = True
@@ -38,7 +44,14 @@ with open('loras.json', 'r') as f:
     
 # Initialize the base model
 #base_model = "stabilityai/stable-diffusion-3.5-large"
-pipe = AutoPipelineForText2Image.from_pretrained("ariG23498/sd-3.5-merged", torch_dtype=torch.bfloat16)
+# Initialize the base model
+dtype = torch.bfloat16
+base_model = "ariG23498/sd-3.5-merged"
+pipe = AutoPipelineForText2Image.from_pretrained(base_model, torch_dtype=dtype).to("cuda")
+#pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
+torch.cuda.empty_cache()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 model_id = ("zer0int/LongCLIP-GmP-ViT-L-14")
 config = CLIPConfig.from_pretrained(model_id)