app.py

import spaces
import gradio as gr
import numpy as np
import random
import torch
from PIL import Image
import re
import paramiko
import urllib
import time
import os
import datetime

from models.transformer_sd3 import SD3Transformer2DModel
from diffusers import StableDiffusion3Pipeline
#from transformers import CLIPTextModelWithProjection, T5EncoderModel
from transformers import CLIPTokenizer, T5TokenizerFast
#from diffusers import SD3Transformer2DModel, AutoencoderKL
#from models.transformer_sd3 import SD3Transformer2DModel
from pipeline_stable_diffusion_3_ipa import StableDiffusion3Pipeline

from image_gen_aux import UpscaleWithModel
from huggingface_hub import hf_hub_download

FTP_HOST = '1ink.us'
FTP_USER = 'ford442'
FTP_PASS = os.getenv("FTP_PASS")
FTP_DIR = '1ink.us/stable_diff'

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_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-large-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

transformer = SD3Transformer2DModel.from_pretrained(
    model_path, subfolder="transformer", torch_dtype=torch.bfloat16
)

pipe = StableDiffusion3Pipeline.from_pretrained(
    #"stabilityai  #  stable-diffusion-3.5-large",
    "ford442/stable-diffusion-3.5-large-bf16",
#    vae=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", use_safetensors=True, subfolder='vae',token=True),
     #scheduler = FlowMatchHeunDiscreteScheduler.from_pretrained('ford442/stable-diffusion-3.5-large-bf16', subfolder='scheduler',token=True),
    # text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
  #  text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
  #  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),
    tokenizer_3=T5TokenizerFast.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", use_fast=True, subfolder="tokenizer_3", token=True),
    torch_dtype=torch.bfloat16,
    #use_safetensors=False,
)

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

pipe.to(device)

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

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

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,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    latent_file,  # Add latents file input
    ip_scale,
    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)
    enhanced_prompt = prompt
    enhanced_prompt_2 = prompt
    
    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).convert('RGB')
        print("-- using image file and loading ip-adapter --")
        pipe.init_ipadapter(
        ip_adapter_path=ip_adapter_path, 
        image_encoder_path=image_encoder_path, 
        nb_token=64, 
        )
        print('-- generating image --')
        #with torch.no_grad():
        sd_image = pipe(
            width=width,
            height=height,
            prompt=enhanced_prompt,  # This conversion is fine
            negative_prompt=negative_prompt_1,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator,
            clip_image=sd_image_a,
            ipadapter_scale=ip_scale,
        ).images[0]
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        rv_path = f"sd35_{timestamp}.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 --')
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        #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"sd35l_{timestamp}.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)
    #  pipe.unet.to('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 --')
    downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
    upscale_path = f"sd35l_upscale_{seed}.png"
    downscale2.save(upscale_path,optimize=False,compress_level=0)
    upload_to_ftp(upscale_path)
    return sd_image, 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;
}
"""

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="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,
            )
            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
            ip_scale = gr.Slider(
                label="Image Prompt Scale",
                minimum=0.0,
                maximum=2.0,
                step=0.01,
                value=0.5,
            )            
            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)"
            )
            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,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            latent_file,  # Add latent_file to the inputs
            ip_scale,
        ],
        outputs=[result, expanded_prompt_output],
        )

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