Update app.py

app.py

@@ -9,7 +9,7 @@ import numpy as np
 import random
 
 import torch
-from diffusers import StableDiffusion3Pipeline
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
 from transformers import CLIPTextModelWithProjection, T5EncoderModel
 from transformers import CLIPTokenizer, T5TokenizerFast
 
@@ -22,7 +22,6 @@ from huggingface_hub import hf_hub_download
 import datetime
 import cyper
 
-from diffusers import AutoencoderKL
 #from models.transformer_sd3 import SD3Transformer2DModel
 #from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
 
@@ -82,6 +81,7 @@ pipe = StableDiffusion3Pipeline.from_pretrained(
   #  text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True),
     #tokenizer=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer", token=True),
     #tokenizer_2=CLIPTokenizer.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=True, subfolder="tokenizer_2", token=True),
+    transformer=None,
     tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", add_prefix_space=False, use_fast=True, subfolder="tokenizer_3", token=True),
     #torch_dtype=torch.bfloat16,
     #use_safetensors=False,
@@ -89,12 +89,13 @@ pipe = StableDiffusion3Pipeline.from_pretrained(
 text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
 text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
 text_encoder_3=T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
-    
+ll_transformer=SD3Transformer2DModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='transformer',token=True).to(torch.device("cuda:0"), dtype=torch.bfloat16)
+pipe.transformer=ll_transformer
 pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")
 
 pipe.to(device=device, dtype=torch.bfloat16)
 #pipe.to(device)
-pipe.vae=vaeX.to('cpu')
+#pipe.vae=vaeX.to('cpu')
 upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cpu'))
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -113,7 +114,8 @@ def infer_30(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipe.vae.to('cpu')
+    pipe.vae=vaeX.to('cpu')
+    pipe.transformer=ll_transformer
     pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
@@ -163,7 +165,8 @@ def infer_60(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipe.vae.to('cpu')
+    pipe.vae=vaeX.to('cpu')
+    pipe.transformer=ll_transformer
     pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
@@ -213,7 +216,8 @@ def infer_90(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipe.vae.to('cpu')
+    pipe.vae=vaeX.to('cpu')
+    pipe.transformer=ll_transformer
     pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
@@ -263,7 +267,8 @@ def infer_100(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
-    pipe.vae.to('cpu')
+    pipe.vae=vaeX.to('cpu')
+    pipe.transformer=ll_transformer
     pipe.text_encoder=text_encoder #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_2=text_encoder_2 #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True).to(device=device, dtype=torch.bfloat16)
     pipe.text_encoder_3=text_encoder_3 #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',token=True).to(device=device, dtype=torch.bfloat16)
@@ -307,7 +312,7 @@ body{background-color: blue;}
 
 with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # StableDiffusion 3.5 Large with UltraReal lora")
+        gr.Markdown(" # StableDiffusion 3.5 Large with UltraReal lora test")
         expanded_prompt_output = gr.Textbox(label="Prompt", lines=1)  # Add this line
         with gr.Row():
             prompt = gr.Text(