Update app.py

app.py

@@ -19,6 +19,10 @@ import glob
 from datetime import datetime
 from PIL import Image
 
+from onediffusion.diffusion.pipelines.onediffusion import OneDiffusionPipeline
+from onediffusion.models.denoiser.nextdit import NextDiT
+from onediffusion.dataset.utils import get_closest_ratio, ASPECT_RATIO_512
+
 #import os
 #cache_dir = '/workspace/hf_cache'
 
@@ -33,44 +37,10 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # Model configurations
 MODEL_CONFIGS = {
-    "FLUX": {
-        "repo_id": "black-forest-labs/FLUX.1-dev",
-        "pipeline_class": FluxPipeline,
-        # "cache_dir" : cache_dir
-    },
-    "Stable Diffusion 3.5": {
-        "repo_id": "stabilityai/stable-diffusion-3.5-large",
-        "pipeline_class": StableDiffusion3Pipeline,
-        # "cache_dir" : cache_dir
-    },
-    "PixArt": {
-        "repo_id": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
-        "pipeline_class": PixArtSigmaPipeline,
-        # "cache_dir" : cache_dir
-    },
-    "SANA": {
-        "repo_id": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
-        "pipeline_class": SanaPipeline,
-        # "cache_dir" : cache_dir
-    },
-    "AuraFlow": {
-        "repo_id": "fal/AuraFlow",
-        "pipeline_class": AuraFlowPipeline,
-        # "cache_dir" : cache_dir
-    },
-    "Kandinsky": {
-        "repo_id": "kandinsky-community/kandinsky-3",
-        "pipeline_class": Kandinsky3Pipeline,
-        #"cache_dir" : cache_dir
-    },
-    "Hunyuan": {
-        "repo_id": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
-        "pipeline_class": HunyuanDiTPipeline,
-        # "cache_dir" : cache_dir
-    },
-    "Lumina": {
-        "repo_id": "Alpha-VLLM/Lumina-Next-SFT-diffusers",
-        "pipeline_class": LuminaText2ImgPipeline,
+
+    "OneDiffusion": {
+        "repo_id": "lehduong/OneDiffusion",
+        "pipeline_class": OneDiffusionPipeline,
         # "cache_dir" : cache_dir
     }
 
@@ -227,6 +197,9 @@ def generate_image(
             print(f"Generating image with {model_name}...")
             # progress(0.3, desc=f"Generating image with {model_name}...")
 
+            if model_name == "OneDiffusion":
+                prompt = "[[text2image]] " + prompt
+            
             image = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,