Update app.py

app.py

@@ -1,7 +1,8 @@
+import os, gc
 import gradio as gr
 import numpy as np
 import random
-
+from transformers import AutoTokenizer, AutoFeatureExtractor
 import spaces
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
 import torch
@@ -15,9 +16,32 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.bfloat16
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler", shift=5, use_fast=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_repo_id,
+    trust_remote_code=True,
+    use_fast=True
+)
+feature_extractor = AutoFeatureExtractor.from_pretrained(
+    model_repo_id,
+    trust_remote_code=True
+)
+
+# 3) Dispatch & load in FP16 with offloading
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo_id,
+    scheduler=FlowMatchEulerDiscreteScheduler.from_pretrained(
+        model_repo_id,
+        subfolder="scheduler",
+        shift=5,
+        use_safetensors=True
+    ),
+    tokenizer=tokenizer,
+    feature_extractor=feature_extractor,
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True,
+    device_map="auto",          # automatically spreads submodules CPU/GPU
+    offload_folder="offload"    # where to spill CPU-offloaded weights
+)
 
 pipe = pipe.to(device)