Update app.py

app.py

@@ -6,18 +6,46 @@ import json
 import torch
 import spaces
 
-from diffusers import Lumina2Text2ImgPipeline
+from diffusers.pipelines import Lumina2Text2ImgPipeline
+from diffusers.models.transformers.transformer_lumina2 import Lumina2Transformer2DModel
+
+from diffusers import (
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler
+)
+from diffusers.loaders.single_file_utils import (
+    convert_sd3_transformer_checkpoint_to_diffusers,
+)
+from transformers import (
+    Gemma2Model,
+    GemmaTokenizer
+)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Alpha-VLLM/Lumina-Image-2.0"
 
-
 if torch.cuda.is_available():
     torch_dtype = torch.bfloat16
 else:
     torch_dtype = torch.float32
 
-pipe = Lumina2Text2ImgPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+###
+
+vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
+text_encoder = Gemma2Model.from_pretrained(model_repo_id, subfolder="text_encoder")
+transformer = Lumina2Transformer2DModel.from_pretrained(model_repo_id, subfolder="transformer")
+tokenizer = GemmaTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler")
+
+###
+pipe = Lumina2Text2ImgPipeline(
+    vae=vae,
+    text_encoder=text_encoder,
+    transformer=transformer,
+    tokenizer=tokenizer,
+    scheduler=scheduler,
+)
+pipe.to(device, torch_dtype)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1536