Update app.py

app.py

@@ -244,73 +244,41 @@ loras = [
     # add new
 ]
 
-# Configurações iniciais
-dtype = torch.float16
+# Initialize the base model
+use_auth_token=True
+dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# URLs e caminhos locais
-config_url = "https://huggingface.co/spaces/vcollos/family/resolve/main/config.json"
-local_config_path = "./config.json"
-vae_weights_url = "https://huggingface.co/spaces/vcollos/family/resolve/main/ae.safetensors"
-vae_weights_path = "./ae.safetensors"
-
-# Função para baixar arquivos
-def download_file(url, local_path):
-    if not os.path.exists(local_path):
-        print(f"Baixando {local_path} de {url}...")
-        response = requests.get(url)
-        response.raise_for_status()
-        with open(local_path, "wb") as f:
-            f.write(response.content)
-        print(f"Arquivo salvo em {local_path}")
-    else:
-        print(f"{local_path} já existe.")
-
-# Baixar arquivos necessários
-download_file(config_url, local_config_path)
-download_file(vae_weights_url, vae_weights_path)
-
-# Carregar o Autoencoder com config.json local
-try:
-    print("Carregando o Autoencoder...")
-    good_vae = AutoencoderKL.from_pretrained(
-        pretrained_model_name_or_path=local_config_path,  # Config local
-        filename=vae_weights_path,  # Pesos locais
-        torch_dtype=dtype
-    ).to(device)
-    print("Autoencoder carregado com sucesso!")
-except Exception as e:
-    print(f"Erro ao carregar o Autoencoder: {e}")
-    good_vae = None
-
-# Configurar o pipeline principal manualmente
-try:
-    print("Configurando o pipeline principal...")
-    pipe = DiffusionPipeline.from_pretrained(
-        pretrained_model_name_or_path=local_config_path,  # Config local
-        torch_dtype=dtype,
-        vae=good_vae
-    ).to(device)
-    print("Pipeline principal configurado com sucesso!")
-except Exception as e:
-    print(f"Erro ao configurar o pipeline principal: {e}")
-    pipe = None
-
-# Configurar o pipeline Image-to-Image
-try:
-    print("Configurando o pipeline Image-to-Image...")
-    pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
-        pretrained_model_or_path=local_config_path,
-        vae=good_vae,
-        torch_dtype=dtype
-    )
-    print("Pipeline Image-to-Image configurado com sucesso!")
-except Exception as e:
-    print(f"Erro ao configurar o pipeline Image-to-Image: {e}")
-
+base_model = "black-forest-labs/FLUX.1-dev"
+
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(
+    base_model,
+    vae=good_vae,
+    transformer=pipe.transformer,
+    text_encoder=pipe.text_encoder,
+    tokenizer=pipe.tokenizer,
+    text_encoder_2=pipe.text_encoder_2,
+    tokenizer_2=pipe.tokenizer_2,
+    torch_dtype=dtype
+)
 
 
-    
+#TAEF1 is very tiny autoencoder which uses the same "latent API" as FLUX.1's VAE. FLUX.1 is useful for real-time previewing of the FLUX.1 generation process.#
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
+pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
+                                                      vae=good_vae,
+                                                      transformer=pipe.transformer,
+                                                      text_encoder=pipe.text_encoder,
+                                                      tokenizer=pipe.tokenizer,
+                                                      text_encoder_2=pipe.text_encoder_2,
+                                                      tokenizer_2=pipe.tokenizer_2,
+                                                      torch_dtype=dtype
+                                                     )
 
 MAX_SEED = 2**32-1