Update app.py

app.py

@@ -8,8 +8,8 @@ import random
 import torch
 from diffusers import StableDiffusion3Pipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler 
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
-from transformers import pipeline
+#from threading import Thread
+#from transformers import pipeline
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 import re
 import paramiko
@@ -118,6 +118,7 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    latent_file,  # Add latents file input
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = random.randint(0, MAX_SEED)
@@ -154,6 +155,9 @@ def infer(
     enhanced_prompt = filter_text(enhanced_prompt,prompt)
     print('-- filtered prompt --')
     print(enhanced_prompt)
+    if latent_file:  # Check if a latent file is provided
+        sd_image_a = torch.load(latent_file.name)  # Load the latent
+        print("-- using latent file --")
     print('-- generating image --')
     with torch.no_grad():
       sd_image = pipe(
@@ -169,6 +173,16 @@ def infer(
     image_path = f"sd35m_{seed}.png"
     sd_image.save(image_path,optimize=False,compress_level=0)
     upload_to_ftp(image_path)
+    
+    # Convert the generated image to a tensor
+    generated_image_tensor = torch.tensor([np.array(sd_image).transpose(2, 0, 1)]).to('cuda') / 255.0
+    # Encode the generated image into latents
+    with torch.no_grad():
+        generated_latents = vae.encode(generated_image_tensor).latent_dist.sample().mul_(0.18215)
+    latent_path = f"sd35m_{seed}.pt"
+    # Save the latents to a .pt file
+    torch.save(generated_latents, latent_path)
+    upload_to_ftp(latent_path)
     #refiner.scheduler.set_timesteps(num_inference_steps,device)
     refine = refiner(
             prompt=f"{prompt}, high quality masterpiece, complex details",
@@ -178,7 +192,7 @@ def infer(
             image=sd_image,
             generator=generator,
     ).images[0]  
-    refine_path = f"refine_{seed}.png"
+    refine_path = f"sd35m_refine_{seed}.png"
     refine.save(refine_path,optimize=False,compress_level=0)
     upload_to_ftp(refine_path)
     return refine, seed, refine_path, enhanced_prompt
@@ -227,7 +241,7 @@ def repeat_infer(
     return result, seed, image_path, enhanced_prompt
 
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(theme=gr.themes.Origin()) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # Text-to-Text-to-Image StableDiffusion 3.5 Medium (with refine)")
         expanded_prompt_output = gr.Textbox(label="Expanded Prompt", lines=5)  # Add this line
@@ -244,6 +258,7 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0, variant="primary")
         result = gr.Image(label="Result", show_label=False)
         with gr.Accordion("Advanced Settings", open=False):
+            latent_file = gr.File(label="Latents File (optional)")  # Add latents file input
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -310,6 +325,7 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            latent_file,  # Add latent_file to the inputs
         ],
         outputs=[result, seed, image_path_output, expanded_prompt_output],
         )