Update app.py

app.py

@@ -7,44 +7,17 @@ import torch
 import spaces
 
 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
-)
 
 default_system_prompt = "You are an assistant designed to generate superior images with the superior degree of image-text alignment based on textual prompts or user prompts."
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Alpha-VLLM/Lumina-Image-2.0"
-transformer_repo_id = "benjamin-paine/Lumina-Image-2.0" # Temporarily fixed, change when main repo gets updated
+
 if torch.cuda.is_available():
     torch_dtype = torch.bfloat16
 else:
     torch_dtype = torch.float32
     
-###
-transformer = Lumina2Transformer2DModel.from_pretrained(transformer_repo_id, subfolder="transformer")
-vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
-text_encoder = Gemma2Model.from_pretrained(model_repo_id, subfolder="text_encoder")
-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 = Lumina2Text2ImgPipeline.from_pretrained(model_repo_id)
 pipe.to(device, torch_dtype)
 
 MAX_SEED = np.iinfo(np.int32).max