app_v5.py

# app_v4.py
import gradio as gr
import torch
from gradio_client import Client, handle_file
import spaces
import os
import datetime
import io
import numpy as np
import moondream as md
from transformers import T5EncoderModel
from diffusers import FluxControlNetPipeline, FluxControlNetInpaintPipeline, FluxTransformer2DModel
from diffusers.utils import load_image
from PIL import Image
from threading import Thread
from typing import Generator
from huggingface_hub import CommitScheduler, HfApi
from debug import log_params, scheduler, save_image
from huggingface_hub.utils._runtime import dump_environment_info
import logging

#############################
os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
#############################

presets = {
    "Strict Upscale": {
        "scale": 1.0,
        "steps": 8,
        "controlnet_conditioning_scale": 0.75,
        "guidance_scale": 4.0,
        "guidance_end": 0.9
    },
    "Creative Upscale": {
        "scale": 2.0,
        "steps": 6,
        "controlnet_conditioning_scale": 0.42,
        "guidance_scale": 3.0,
        "guidance_end": 0.5
    },
    "High Detail Upscale": {
        "scale": 1.25,
        "steps": 10,
        "controlnet_conditioning_scale": 0.9,
        "guidance_scale": 10.0,
        "guidance_end": 0.9
    }
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = 1000000

huggingface_token = os.getenv("HUGGINFACE_TOKEN")
md_api_key = os.getenv("MD_KEY")
model = md.vl(api_key=md_api_key)

try:
    # Set max memory usage for ZeroGPU
    torch.cuda.set_per_process_memory_fraction(1.0)
    torch.set_float32_matmul_precision("high")
except Exception as e:
    print(f"Error setting memory usage: {e}")

text_encoder_2_unquant = T5EncoderModel.from_pretrained(
    "LPX55/FLUX.1-merged_lightning_v2",
    subfolder="text_encoder_2",
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)
transformer = FluxTransformer2DModel.from_pretrained(
        "LPX55/FLUX.1-merged_lightning_v2", subfolder='transformer', torch_dytpe=torch.bfloat16
    )
pipe_upscaler = FluxControlNetPipeline.from_pretrained(
    "LPX55/FLUX.1M-8step_upscaler-cnet",
    torch_dtype=torch.bfloat16,
    text_encoder_2=text_encoder_2_unquant,
    token=huggingface_token
)
pipe_upscaler.to("cuda")

controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16)

pipe = FluxControlNetInpaintPipeline.from_pretrained(
    "LPX55/FLUX.1-merged_lightning_v2",
    controlnet=controlnet,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)
pipe.to("cuda")
pipe.transformer.to(torch.bfloat16)
pipe.controlnet.to(torch.bfloat16)

try:
    dump_environment_info()
except Exception as e:
    print(f"Failed to dump env info: {e}")
    
def get_image_dimensions(image: Image) -> str:
    width, height = image.size
    return f"Original Image Dimensions: {width}x{height}"
    
def resize_image_to_max_side(image: Image, max_side_length=1024) -> Image:
    width, height = image.size
    ratio = min(max_side_length / width, max_side_length / height)
    new_size = (int(width * ratio), int(height * ratio))
    resized_image = image.resize(new_size, Image.BILINEAR)
    return resized_image


def combine_caption_focus(caption, focus):
    try:
        if caption is None:
            caption = ""
        if focus is None:
            focus = "highly detailed photo, raw photography."
        return (str(caption) + "\n\n" + str(focus)).strip()
    except Exception as e:
        print(f"Error combining caption and focus: {e}")
        return "highly detailed photo, raw photography."
def generate_caption(control_image):
    try:
        if control_image is None:
            return "Waiting for control image...", "Original Image Dimensions: N/A"
           
        # Get original dimensions
        original_dimensions = get_image_dimensions(control_image)
        
        # Resize the image to a maximum longest side of 1024 pixels
        control_image = resize_image_to_max_side(control_image, max_side_length=1024)

        # Generate a detailed caption
        mcaption = model.caption(control_image, length="short")
        detailed_caption = mcaption["caption"]
        print(f"Detailed caption: {detailed_caption}")
        
        return detailed_caption, original_dimensions
    except Exception as e:
        print(f"Error generating caption: {e}")
        return "A detailed photograph", "Original Image Dimensions: N/A"

def generate_focus(control_image, focus_list):
    try:
        if control_image is None:
            return None, "Original Image Dimensions: N/A"
        if focus_list is None:
            return "", "Original Image Dimensions: N/A"
        
        # Get original dimensions
        original_dimensions = get_image_dimensions(control_image)
        
        # Resize the image to a maximum longest side of 1024 pixels
        control_image = resize_image_to_max_side(control_image, max_side_length=1024)

        # Generate a detailed caption
        focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
        focus_description = focus_query["answer"]
        print(f"Areas of focus: {focus_description}") 
        return focus_description, original_dimensions
    except Exception as e:
        print(f"Error generating focus: {e}")
        return "highly detailed photo, raw photography.", "Original Image Dimensions: N/A"
        

@spaces.GPU(duration=6, progress=gr.Progress(track_tqdm=True))
@torch.no_grad()
def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
    generator = torch.Generator().manual_seed(seed)
    
    # Ensure transparency is preserved
    control_image = control_image.convert("RGBA")
    
    # Resize the image to a maximum longest side of 1024 pixels
    control_image = resize_image_to_max_side(control_image, max_side_length=1024)
    
    w, h = control_image.size
    # Crop to nearest multiple of 32
    w = w - w % 32
    h = h - h % 32
    # Corrected resizing code
    control_image = control_image.resize((w, h), resample=2)
    
    print(f"Resized image dimensions: {control_image.size[0]}x{control_image.size[1]}")
    print(f"PromptLog: {repr(prompt)}")
    
    # Convert image to RGB for processing, but keep alpha channel for transparency
    control_image_rgb = control_image.convert("RGB")
    control_image_alpha = control_image.split()[-1]
    
    # Convert alpha channel to a mask (transparent = white, opaque = black)
    # White corresponds to 1 (to be inpainted), black corresponds to 0 (to be preserved)
    # Convert alpha to numpy array for processing
    alpha_array = np.array(control_image_alpha)
    # Create binary mask (1 for transparent, 0 for opaque)
    mask = (alpha_array > 128).astype(np.float32)  # 1 for transparent (to be inpainted), 0 for opaque
    
    # Optional: Visualize the mask (for debugging purposes)
    # mask_image = Image.fromarray((mask * 255).astype(np.uint8))
    # mask_image.show()
    
    with torch.inference_mode():
        image = pipe(
            image=control_image_rgb,
            generator=generator,
            prompt=prompt,
            control_image=control_image_rgb,
            mask_image=mask,  # Pass the numpy array as the mask
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            height=control_image.size[1],
            width=control_image.size[0],
            control_guidance_start=0.0,
            control_guidance_end=guidance_end,
        ).images[0]
        
        # Reapply the alpha channel to the generated image
        image = image.convert("RGBA")
        image.putalpha(control_image_alpha)
        
        return image
    
def update_parameters(preset):
    if preset in presets:
        params = presets[preset]
        return (
            params["scale"],
            params["steps"],
            params["controlnet_conditioning_scale"],
            params["guidance_scale"],
            params["guidance_end"]
        )
    else:
        # Default values if preset is not found
        return 1.0, 8, 0.6, 3.5, 1.0

def process_image(control_image, user_prompt, system_prompt, scale, steps, 
                controlnet_conditioning_scale, guidance_scale, seed, 
                guidance_end, temperature, max_new_tokens, log_prompt):
    # Initialize with empty caption
    final_prompt = user_prompt.strip()
    # If no user prompt provided, generate a caption first
    if not final_prompt:
        # Generate a detailed caption

        mcaption = model.caption(control_image, length="normal")
        detailed_caption = mcaption["caption"]
        final_prompt = detailed_caption
        yield f"Using caption: {final_prompt}", None, final_prompt

    yield f"Generating with: {final_prompt}", None, final_prompt

    # Generate the image
    try:
        image = generate_image(
            prompt=final_prompt,
            scale=scale,
            steps=steps,
            control_image=control_image,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            guidance_scale=guidance_scale,
            seed=seed,
            guidance_end=guidance_end
        )
        try:
            # Ensure the image is saved with transparency
            with io.BytesIO() as output:
                image.save(output, format="PNG")
                debug_img = Image.open(output).convert("RGBA")
                save_image("/tmp/" + str(seed) + "output.png", debug_img)
        except Exception as e:
            print("Error 160: " + str(e))
        log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
        yield f"Completed! Used prompt: {final_prompt}", image, final_prompt
    except Exception as e:
        print("Error: " + str(e))
        yield f"Error: {str(e)}", None, None

with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
    gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
    # status_box = gr.Markdown("🔄 Warming up...")
    
    with gr.Row():
        with gr.Accordion():
            control_image = gr.Image(type="pil", label="Control Image", show_label=False)
        with gr.Accordion():
            generated_image = gr.Image(type="pil", label="Generated Image", format="png", show_label=False)
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(lines=4, info="Enter your prompt here or wait for auto-generation...", label="Image Description")
            focus = gr.Textbox(label="Area(s) of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.", value="clothing material, textures, ethnicity")
            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.1)
            with gr.Row():
                generate_button = gr.Button("Generate Image", variant="primary")
                caption_button = gr.Button("Generate Caption", variant="secondary")
        with gr.Column(scale=1):
            with gr.Row():
                preset_choices = list(presets.keys())
                preset_radio = gr.Radio(choices=preset_choices, label="Select Preset", value=preset_choices[0])
            seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
            steps = gr.Slider(2, 16, value=8, label="Steps", step=1)
            controlnet_conditioning_scale = gr.Slider(0, 1, value=0.6, label="ControlNet Scale")
            guidance_scale = gr.Slider(1, 30, value=3.5, label="Guidance Scale")
            guidance_end = gr.Slider(0, 1, value=1.0, label="Guidance End")
            original_dimensions = gr.Markdown(value="Original Image Dimensions: N/A")  # New output for dimensions

    with gr.Row():
        with gr.Accordion("Auto-Caption settings", open=False, visible=False):
            system_prompt = gr.Textbox(
                lines=4, 
                value="Write a straightforward caption for this image. Begin with the main subject and medium. Mention pivotal elements—people, objects, scenery—using confident, definite language. Focus on concrete details like color, shape, texture, and spatial relationships. Show how elements interact. Omit mood and speculative wording. If text is present, quote it exactly. Note any watermarks, signatures, or compression artifacts. Never mention what's absent, resolution, or unobservable details. Vary your sentence structure and keep the description concise, without starting with 'This image is…' or similar phrasing.",
                label="System Prompt for Captioning",
                visible=False  # Changed to visible
            )
            temperature_slider = gr.Slider(
                minimum=0.0, maximum=2.0, value=0.6, step=0.05,
                label="Temperature",
                info="Higher values make the output more random, lower values make it more deterministic.",
                visible=False  # Changed to visible
            )
            max_tokens_slider = gr.Slider(
                minimum=1, maximum=2048, value=368, step=1,
                label="Max New Tokens",
                info="Maximum number of tokens to generate. The model will stop generating if it reaches this limit.",
                visible=False  # Changed to visible
            )
        log_prompt = gr.Checkbox(value=True, label="Log", visible=False)  # Changed to visible
    
    gr.Markdown("**Tips:** 8 steps is all you need! Incredibly powerful tool, usage instructions coming soon.")
    with gr.Accordion("Auth for those Getting ZeroGPU errors.", open=False, elem_id="zgpu"):
        msg1 = gr.Markdown()
        try_btn = gr.LoginButton()
        
        # sus = ['x-zerogpu-token', 'x-zerogpu-uuid', 'x-proxied-host', 'x-proxied-path', 'x-proxied-replica', 'x-request-id', 'x-ip-token']
        # x_ip_token = request.headers['X-IP-TOKEN']
        # print(str(x_ip_token))
        # client = Client("LPX55/zerogpu-experiments", hf_token=huggingface_token, headers={"x-ip-token": x_ip_token})
        # cresult = client.predict(
        #         n=3,
        #         api_name="/predict"
        # )

    caption_state = gr.State()
    focus_state = gr.State()
    log_state = gr.State()
    generate_button.click(
        fn=process_image,
        inputs=[
            control_image, prompt, system_prompt, scale, steps, 
            controlnet_conditioning_scale, guidance_scale, seed, 
            guidance_end, temperature_slider, max_tokens_slider, log_prompt
        ],
        outputs=[log_state, generated_image, prompt]
    )
    control_image.upload(
        generate_caption,
        inputs=[control_image],
        outputs=[caption_state, original_dimensions]
    ).then(
        generate_focus,
        inputs=[control_image, focus],
        outputs=[focus_state, original_dimensions]
    ).then(
        combine_caption_focus,
        inputs=[caption_state, focus_state],
        outputs=[prompt]
    )
    caption_button.click(
        fn=generate_caption,
        inputs=[control_image],
        outputs=[prompt, original_dimensions]
    ).then(
        generate_focus,
        inputs=[control_image, focus],
        outputs=[focus_state, original_dimensions]
    ).then(
        combine_caption_focus,
        inputs=[caption_state, focus_state],
        outputs=[prompt]
    )
    preset_radio.change(
        fn=update_parameters,
        inputs=[preset_radio],
        outputs=[scale, steps, controlnet_conditioning_scale, guidance_scale, guidance_end]
    )
    def hello(profile: gr.OAuthProfile | None) -> str:
        if profile is None:
            return "Hello guest! There is a bug with HF ZeroGPUs that are afffecting some usage on certain spaces. Testing out some possible solutions."
        return f"You are logged in as {profile.name}. If you run into incorrect messages about ZeroGPU runtime credits being out, PLEASE give me a heads up so I can investigate further."

    demo.load(hello, inputs=None, outputs=msg1)
demo.queue().launch(show_error=True)