app.py

import spaces
import gradio as gr
import numpy as np
#import tensorrt as trt
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 transformers import T5Tokenizer, T5ForConditionalGeneration
import re
import paramiko
import urllib
import time
import os
from image_gen_aux import UpscaleWithModel
from huggingface_hub import hf_hub_download
#from models.transformer_sd3 import SD3Transformer2DModel
#from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline
from PIL import Image

FTP_HOST = "1ink.us"
FTP_USER = "ford442"
FTP_PASS = "GoogleBez12!"
FTP_DIR = "1ink.us/stable_diff/"  # Remote directory on FTP server

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = False
#torch.backends.cudnn.benchmark = False
torch.backends.cuda.preferred_blas_library="cublas"
#torch.backends.cuda.preferred_linalg_library="cusolver"

hftoken = os.getenv("HF_AUTH_TOKEN") 

image_encoder_path = "google/siglip-so400m-patch14-384"
ipadapter_path = hf_hub_download(repo_id="InstantX/SD3.5-Large-IP-Adapter", filename="ip-adapter.bin")
model_path = 'ford442/stable-diffusion-3.5-medium-bf16'

def upload_to_ftp(filename):
    try:
        transport = paramiko.Transport((FTP_HOST, 22))
        destination_path=FTP_DIR+filename
        transport.connect(username = FTP_USER, password = FTP_PASS)
        sftp = paramiko.SFTPClient.from_transport(transport)
        sftp.put(filename, destination_path)
        sftp.close()
        transport.close()
        print(f"Uploaded {filename} to FTP server")
    except Exception as e:
        print(f"FTP upload error: {e}")

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch_dtype = torch.bfloat16

checkpoint = "microsoft/Phi-3.5-mini-instruct"
#vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
#vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
#vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #, device_map='cpu') #.to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)

pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-large-bf16").to(device=device, dtype=torch.bfloat16)
#pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(torch.device("cuda:0"))
#pipe = StableDiffusion3Pipeline.from_pretrained("ford442/RealVis_Medium_1.0b_bf16", torch_dtype=torch.bfloat16)
#pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')

# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")

#pipe.scheduler.config.requires_aesthetics_score = False
#pipe.enable_model_cpu_offload()
#pipe.to(device)
#pipe.to(device=device, dtype=torch.bfloat16)
#pipe = torch.compile(pipe)
# pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear")

#refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16",vae = vaeXL, requires_aesthetics_score=True) #.to(torch.bfloat16)
#refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float32, requires_aesthetics_score=True, device_map='balanced')
#refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
#refiner.enable_model_cpu_offload()

#pipe.to(device=device, dtype=torch.bfloat16)

#refiner.scheduler.config.requires_aesthetics_score=False
#refiner.to(device)
#refiner = torch.compile(refiner)
#refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config, beta_schedule="scaled_linear")
#refiner.scheduler = EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)

tokenizer = AutoTokenizer.from_pretrained(checkpoint, add_prefix_space=True)
tokenizer.tokenizer_legacy=False
model = AutoModelForCausalLM.from_pretrained(checkpoint).to('cuda')
#model = torch.compile(model)

upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))

def filter_text(text,phraseC):
  """Filters out the text up to and including 'Rewritten Prompt:'."""
  phrase = "Rewritten Prompt:"
  phraseB = "rewritten text:"
  pattern = f"(.*?){re.escape(phrase)}(.*)"
  patternB = f"(.*?){re.escape(phraseB)}(.*)"
  #  matchB = re.search(patternB, text)
  matchB = re.search(patternB, text, flags=re.DOTALL)
  if matchB:
        filtered_text = matchB.group(2)
        match = re.search(pattern, filtered_text, flags=re.DOTALL)
        if match:
          filtered_text = match.group(2)
          filtered_text = re.sub(phraseC, "", filtered_text, flags=re.DOTALL)  # Replaces the matched pattern with an empty string
          return filtered_text
        else:
          return filtered_text
  else:
        # Handle the case where no match is found
        return text

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 4096

@spaces.GPU(duration=90)
def infer(
    prompt,
    negative_prompt_1,
    negative_prompt_2,
    negative_prompt_3,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    expanded,
    latent_file,  # Add latents file input
    progress=gr.Progress(track_tqdm=True),
):
    upscaler_2.to(torch.device('cpu'))
    torch.set_float32_matmul_precision("highest")
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device='cuda').manual_seed(seed)
    if expanded:
        system_prompt_rewrite = (
        "You are an AI assistant that rewrites image prompts to be more descriptive and detailed."
        )
        user_prompt_rewrite = (
        "Rewrite this prompt to be more descriptive and detailed and only return the rewritten text: "
        )
        user_prompt_rewrite_2 = (
        "Rephrase this scene to have more elaborate details: "
        )
        input_text = f"{system_prompt_rewrite} {user_prompt_rewrite} {prompt}"
        input_text_2 = f"{system_prompt_rewrite} {user_prompt_rewrite_2} {prompt}"
        print("-- got prompt --")
        # Encode the input text and include the attention mask
        encoded_inputs = tokenizer(input_text, return_tensors="pt", return_attention_mask=True)
        encoded_inputs_2 = tokenizer(input_text_2, return_tensors="pt", return_attention_mask=True)
        # Ensure all values are on the correct device
        input_ids = encoded_inputs["input_ids"].to(device)
        input_ids_2 = encoded_inputs_2["input_ids"].to(device)
        attention_mask = encoded_inputs["attention_mask"].to(device)
        attention_mask_2 = encoded_inputs_2["attention_mask"].to(device)
        print("-- tokenize prompt --")
          # Google T5
        #input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=512,
            temperature=0.2,
            top_p=0.9,
            do_sample=True,
        )
        outputs_2 = model.generate(
            input_ids=input_ids_2,
            attention_mask=attention_mask_2,
            max_new_tokens=65,
            temperature=0.2,
            top_p=0.9,
            do_sample=True,
        )
        # Use the encoded tensor 'text_inputs' here
        enhanced_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
        enhanced_prompt_2 = tokenizer.decode(outputs_2[0], skip_special_tokens=True)
        print('-- generated prompt --')
        enhanced_prompt = filter_text(enhanced_prompt,prompt)
        enhanced_prompt_2 = filter_text(enhanced_prompt_2,prompt)
        print('-- filtered prompt --')
        print(enhanced_prompt)
        print('-- filtered prompt 2 --')
        print(enhanced_prompt_2)
    else:
        enhanced_prompt = prompt
        enhanced_prompt_2 = prompt
    model.to('cpu')
    if latent_file:  # Check if a latent file is provided
      #  initial_latents = pipe.prepare_latents(
      #      batch_size=1,
      #      num_channels_latents=pipe.transformer.in_channels,
      #      height=pipe.transformer.config.sample_size[0],
       #     width=pipe.transformer.config.sample_size[1],
      #      dtype=pipe.transformer.dtype,
      #      device=pipe.device,
      #      generator=generator,
      #  )
        sd_image_a = Image.open(latent_file.name)
        print("-- using image file --")
        print('-- generating image --')
        #with torch.no_grad():
        sd_image = pipe(
            prompt=enhanced_prompt,  # This conversion is fine
            negative_prompt=negative_prompt_1,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            latents=sd_image_a,
            generator=generator
        ).images[0]
        rv_path = f"sd35_{seed}.png"
        sd_image[0].save(rv_path,optimize=False,compress_level=0)
        upload_to_ftp(rv_path)
    else:
        print('-- generating image --')
        #with torch.no_grad():
        sd_image = pipe(
            prompt=prompt,  # This conversion is fine
            prompt_2=enhanced_prompt_2,
            prompt_3=enhanced_prompt,
            negative_prompt=negative_prompt_1,
            negative_prompt_2=negative_prompt_2,
            negative_prompt_3=negative_prompt_3,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
         #   latents=None,
          #  output_type='latent',
            generator=generator,
            max_sequence_length=512
        ).images[0]
        print('-- got image --')
        #sd35_image = pipe.vae.decode(sd_image / 0.18215).sample
       # sd35_image = sd35_image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
       # sd35_image = (sd35_image * 255).round().astype("uint8")
       # image_pil = Image.fromarray(sd35_image[0])
      #  sd35_path = f"sd35_{seed}.png"
       # image_pil.save(sd35_path,optimize=False,compress_level=0)
       # upload_to_ftp(sd35_path)
        sd35_path = f"sd35_{seed}.png"
        sd_image.save(sd35_path,optimize=False,compress_level=0)
        upload_to_ftp(sd35_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 = pipe.vae.encode(generated_image_tensor.to(torch.bfloat16)).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)
    '''
    pipe.to(torch.device('cpu'))
    refiner.to(device=device, dtype=torch.bfloat16)
    refine = refiner(
            prompt=f"{enhanced_prompt_2}, high quality masterpiece, complex details",
            negative_prompt = negative_prompt_1,
            negative_prompt_2 = negative_prompt_2,
            guidance_scale=7.5,
            num_inference_steps=num_inference_steps,
            image=sd_image,
            generator=generator,
    ).images[0]
    refine_path = f"sd35_refine_{seed}.png"
    refine.save(refine_path,optimize=False,compress_level=0)
    upload_to_ftp(refine_path)
    refiner.to(torch.device('cpu'))
    '''
    upscaler_2.to(torch.device('cuda'))
    with torch.no_grad():
        upscale2 = upscaler_2(sd_image, tiling=True, tile_width=256, tile_height=256)
    print('-- got upscaled image --')
    upscaler_2.to(torch.device('cpu'))
    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
    upscale_path = f"sd35_upscale_{seed}.png"
    downscale2.save(upscale_path,optimize=False,compress_level=0)
    upload_to_ftp(upscale_path)
    return sd_image, seed, enhanced_prompt

examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
body{
  background-color: blue;
}
"""

def repeat_infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    num_iterations,  # New input for number of iterations
):
    i = 0
    while i < num_iterations:
        time.sleep(700)  # Wait for 10 minutes (600 seconds)
        result, seed, image_path, enhanced_prompt = infer(
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        )
        
        # Optionally, you can add logic here to process the results of each iteration
        # For example, you could display the image, save it with a different name, etc.
        i += 1
    return result, seed, image_path, enhanced_prompt

with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Text-to-Image StableDiffusion 3.5 Large")
        expanded_prompt_output = gr.Textbox(label="Expanded Prompt", lines=5)  # Add this line
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            options = [True, False]
            expanded = gr.Radio(
                show_label=True,
                container=True,
                interactive=True,
                choices=options,
                value=True,
                label="Use expanded prompt: ",
            )
            run_button = gr.Button("Run", scale=0, variant="primary")
        result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("Advanced Settings", open=True):
            latent_file = gr.File(label="Image File (optional)")  # Add latents file input
            negative_prompt_1 = gr.Text(
                label="Negative prompt 1",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=True,
                value="bad anatomy, poorly drawn hands, distorted face, blurry, out of frame, low resolution, grainy, pixelated, disfigured, mutated, extra limbs, bad composition"
            )
            negative_prompt_2 = gr.Text(
                label="Negative prompt 2",
                max_lines=1,
                placeholder="Enter a second negative prompt",
                visible=True,
                value="unrealistic, cartoon, anime, sketch, painting, drawing, illustration, graphic, digital art, render, 3d, blurry, deformed, disfigured, poorly drawn, bad anatomy, mutated, extra limbs, ugly, out of frame, bad composition, low resolution, grainy, pixelated, noisy, oversaturated, undersaturated, (worst quality, low quality:1.3), (bad hands, missing fingers:1.2)"
            )
            negative_prompt_3 = gr.Text(
                label="Negative prompt 3",
                max_lines=1,
                placeholder="Enter a third negative prompt",
                visible=True,
                value="(worst quality, low quality:1.3), (bad anatomy, bad hands, missing fingers, extra digit, fewer digits:1.2), (blurry:1.1), cropped, watermark, text, signature, logo, jpeg artifacts, (ugly, deformed, disfigured:1.2), (poorly drawn:1.2), mutated, extra limbs, (bad proportions, gross proportions:1.2), (malformed limbs, missing arms, missing legs, extra arms, extra legs:1.2), (fused fingers, too many fingers, long neck:1.2), (unnatural body, unnatural pose:1.1), out of frame, (bad composition, poorly composed:1.1), (oversaturated, undersaturated:1.1), (grainy, pixelated:1.1), (low resolution, noisy:1.1), (unrealistic, distorted:1.1), (extra fingers, mutated hands, poorly drawn hands, bad hands:1.3), (missing fingers:1.3)"
            )
            num_iterations = gr.Number(
                value=1000, 
                label="Number of Iterations")
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,  # Replace with defaults that work for your model
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,  # Replace with defaults that work for your model
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=30.0,
                    step=0.1,
                    value=4.2,  # Replace with defaults that work for your model
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=500,
                    step=1,
                    value=220,  # Replace with defaults that work for your model
                )
            gr.Examples(examples=examples, inputs=[prompt])
        gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt_1,
            negative_prompt_2,
            negative_prompt_3,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            expanded,
            latent_file,  # Add latent_file to the inputs
        ],
        outputs=[result, seed, expanded_prompt_output],
        )

if __name__ == "__main__":
    demo.launch()