Update app.py

app.py

@@ -1,12 +1,17 @@
+from dataclasses import dataclass
+from typing import Union, Optional, List, Any, Dict
+
 import gradio as gr
 import numpy as np
 import random
 import spaces
 import torch
-from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from huggingface_hub import hf_hub_download
+
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL, FluxPipeline
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
 
-from model import Flux
+from model import Flux, FluxParams
 
 def calculate_shift(
     image_seq_len,
@@ -174,20 +179,17 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
 @dataclass
 class ModelSpec:
     params: FluxParams
-    ae_params: AutoEncoderParams
-    ckpt_path: str 
-    ae_path: str 
     repo_id: str 
     repo_flow: str 
     repo_ae: str 
     repo_id_ae: str
 
+
 config = ModelSpec(
         repo_id="TencentARC/flux-mini",
         repo_flow="flux-mini.safetensors",
         repo_id_ae="black-forest-labs/FLUX.1-dev",
         repo_ae="ae.safetensors",
-        ckpt_path=os.getenv("FLUX_MINI", None),
         params=FluxParams(
             in_channels=64,
             vec_in_dim=768,
@@ -202,35 +204,33 @@ config = ModelSpec(
             qkv_bias=True,
             guidance_embed=True,
         )
+)
 
 
-def load_flow_model2(device: str = "cuda", hf_download: bool = True):
-    if (
-        and config.repo_id is not None
+def load_flow_model2(config, device: str = "cuda", hf_download: bool = True):
+    if (config.repo_id is not None
         and config.repo_flow is not None
         and hf_download
     ):
-        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow.replace("sft", "safetensors"))
+        ckpt_path = hf_hub_download(config.repo_id, config.repo_flow.replace("sft", "safetensors"))
 
-    model = Flux(params)
+    model = Flux(config.params)
     if ckpt_path is not None:
         sd = load_sft(ckpt_path, device=str(device))
         missing, unexpected = model.load_state_dict(sd, strict=True)
     return model
 
 
-
-
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler").to(device)
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="scheduler")
 vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
 text_encoder = CLIPTextModel.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder").to(device)
-tokenizer = CLIPTokenizer.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="tokenizer").to(device)
+tokenizer = CLIPTokenizer.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="tokenizer")
 text_encoder_2 = T5EncoderModel.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2").to(device)
-tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="tokenizer_2").to(device)
-transformer = load_flow_model2(device)
+tokenizer_2 = T5TokenizerFast.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="tokenizer_2")
+transformer = load_flow_model2(config, device)
 
 pipe = FluxPipeline(
     scheduler,
@@ -238,7 +238,7 @@ pipe = FluxPipeline(
     text_encoder,
     tokenizer,
     text_encoder_2,
-    tokenizer_2
+    tokenizer_2,
     transformer
 )
 torch.cuda.empty_cache()