app.py

import spaces
import os

os.environ["SAFETENSORS_FAST_GPU"] = "1"
os.putenv("HF_HUB_ENABLE_HF_TRANSFER","1")

import gradio as gr
import numpy as np
import random

import torch
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"
torch.set_float32_matmul_precision("highest")

from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, AutoencoderKL
from transformers import CLIPTextModelWithProjection, T5EncoderModel
from transformers import CLIPTokenizer, T5TokenizerFast

import re
import paramiko
import urllib
import time
from image_gen_aux import UpscaleWithModel
from huggingface_hub import hf_hub_download
import datetime
import cyper
from PIL import Image

from accelerate import Accelerator

accelerator = Accelerator(mixed_precision="bf16")

hftoken = os.getenv("HF_AUTH_TOKEN") 

code = r'''
import torch
import paramiko
import os
FTP_HOST = '1ink.us'
FTP_USER = 'ford442'
FTP_PASS = os.getenv("FTP_PASS")
FTP_DIR = '1ink.us/stable_diff/'

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}")
'''

pyx = cyper.inline(code, fast_indexing=True, directives=dict(boundscheck=False, wraparound=False, language_level=3))

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

#vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", use_safetensors=True, subfolder='vae',token=True)
vaeX=AutoencoderKL.from_pretrained("ford442/stable-diffusion-3.5-large-fp32", safety_checker=None, use_safetensors=True, subfolder='vae', low_cpu_mem_usage=False, torch_dtype=torch.float32, token=True)
pipe = StableDiffusion3Pipeline.from_pretrained(
    #"stabilityai  #  stable-diffusion-3.5-large",
    "ford442/stable-diffusion-3.5-large-bf16",
    trust_remote_code=True,
    vae=None,
    #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=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
    # text_encoder=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder', token=True),
    text_encoder_2=None, #CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
  #  text_encoder_2=CLIPTextModelWithProjection.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_2',token=True),
    text_encoder_3=None, #T5EncoderModel.from_pretrained("ford442/stable-diffusion-3.5-large-bf16", subfolder='text_encoder_3',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),
    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,
)
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.eval()
pipe.load_lora_weights("ford442/sdxl-vae-bf16", weight_name="LoRA/UltraReal.safetensors")

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

upscaler_2 = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device('cpu'))

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

MAX_IMAGE_SIZE = 4096

@spaces.GPU(duration=70)
def infer_60(
    prompt,
    negative_prompt_1,
    negative_prompt_2,
    negative_prompt_3,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    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)
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device='cuda').manual_seed(seed)
    print('-- generating image --')
    sd_image = pipe(
            prompt=prompt,
            prompt_2=prompt,
            prompt_3=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,
            generator=generator,
            max_sequence_length=512
    ).images[0]
    print('-- got image --')
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    sd35_path = f"sd35ll_{timestamp}.png"
    sd_image.save(sd35_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(sd35_path)
    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"sd35ll_upscale_{timestamp}.png"
    downscale2.save(upscale_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(upscale_path)
    return sd_image, prompt

@spaces.GPU(duration=100)
def infer_90(
    prompt,
    negative_prompt_1,
    negative_prompt_2,
    negative_prompt_3,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    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)
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device='cuda').manual_seed(seed)
    print('-- generating image --')
    sd_image = pipe(
            prompt=prompt,
            prompt_2=prompt,
            prompt_3=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,
            generator=generator,
            max_sequence_length=512
    ).images[0]
    print('-- got image --')
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    sd35_path = f"sd35ll_{timestamp}.png"
    sd_image.save(sd35_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(sd35_path)
    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"sd35ll_upscale_{timestamp}.png"
    downscale2.save(upscale_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(upscale_path)
    return sd_image, prompt

@spaces.GPU(duration=120)
def infer_110(
    prompt,
    negative_prompt_1,
    negative_prompt_2,
    negative_prompt_3,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    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)
    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device='cuda').manual_seed(seed)
    print('-- generating image --')
    sd_image = pipe(
            prompt=prompt,
            prompt_2=prompt,
            prompt_3=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,
            generator=generator,
            max_sequence_length=512
    ).images[0]
    print('-- got image --')
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    sd35_path = f"sd35ll_{timestamp}.png"
    sd_image.save(sd35_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(sd35_path)
    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"sd35ll_upscale_{timestamp}.png"
    downscale2.save(upscale_path,optimize=False,compress_level=0)
    pyx.upload_to_ftp(upscale_path)
    return sd_image, prompt
    
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(" # 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(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button_60 = gr.Button("Run 60", scale=0, variant="primary")
            run_button_90 = gr.Button("Run 90", scale=0, variant="primary")
            run_button_110 = gr.Button("Run 110", scale=0, variant="primary")
        result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("Advanced Settings", open=True):
            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")
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=768,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=30.0,
                    step=0.1,
                    value=4.2,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=500,
                    step=1,
                    value=50,
                )
        gr.on(
        triggers=[run_button_60.click, prompt.submit],
        fn=infer_60,
        inputs=[
            prompt,
            negative_prompt_1,
            negative_prompt_2,
            negative_prompt_3,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, expanded_prompt_output],
        )
        gr.on(
        triggers=[run_button_90.click, prompt.submit],
        fn=infer_90,
        inputs=[
            prompt,
            negative_prompt_1,
            negative_prompt_2,
            negative_prompt_3,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, expanded_prompt_output],
        )
        gr.on(
        triggers=[run_button_110.click, prompt.submit],
        fn=infer_110,
        inputs=[
            prompt,
            negative_prompt_1,
            negative_prompt_2,
            negative_prompt_3,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, expanded_prompt_output],
        )

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