raw.py

import torch
import spaces
import os
from diffusers.utils import load_image
from diffusers.hooks import apply_group_offloading
from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import T5EncoderModel
from transformers import LlavaForConditionalGeneration, TextIteratorStreamer, AutoProcessor
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from liger_kernel.transformers import apply_liger_kernel_to_llama
from PIL import Image
from threading import Thread
from typing import Generator
from peft import PeftModel, PeftConfig
# from attention_map_diffusers import (
#     attn_maps,
#     init_pipeline,
#     save_attention_maps
# )
import gradio as gr
huggingface_token = os.getenv("HUGGINFACE_TOKEN")
MAX_SEED = 1000000
MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
cap_processor = AutoProcessor.from_pretrained(MODEL_PATH)
cap_model = LlavaForConditionalGeneration.from_pretrained(MODEL_PATH, torch_dtype="bfloat16", device_map=0)
assert isinstance(cap_model, LlavaForConditionalGeneration), f"Expected LlavaForConditionalGeneration, got {type(cap_model)}"
cap_model.eval()
apply_liger_kernel_to_llama(model=cap_model.language_model)
# quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)
# text_encoder_2_8bit = T5EncoderModel.from_pretrained(
#     "LPX55/FLUX.1-merged_uncensored",
#     subfolder="text_encoder_2",
#     quantization_config=quant_config,
#     torch_dtype=torch.bfloat16,
#     token=huggingface_token
# )
text_encoder_2_unquant = T5EncoderModel.from_pretrained(
    "LPX55/FLUX.1-merged_uncensored",
    subfolder="text_encoder_2",
    torch_dtype=torch.bfloat16,
    token=huggingface_token
)
# good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")

# Load pipeline
# controlnet = FluxControlNetModel.from_pretrained(
#     "jasperai/Flux.1-dev-Controlnet-Upscaler",
#     torch_dtype=torch.bfloat16
# )
pipe = FluxControlNetPipeline.from_pretrained(
    "LPX55/FLUX.1M-8step_upscaler-cnet",
    torch_dtype=torch.bfloat16,
    text_encoder_2=text_encoder_2_unquant,
    token=huggingface_token
)
# adapter_id = "alimama-creative/FLUX.1-Turbo-Alpha"
# adapter_id2 = "XLabs-AI/flux-RealismLora"
# adapter_id3 = "enhanceaiteam/Flux-uncensored-v2"

pipe.to("cuda")

# try:
#     pipe.vae.enable_slicing()
# except:
#     print("debug-2")
# try:
#     pipe.vae.enable_tiling()
# except:
#     print("debug-3")

# pipe.load_lora_weights(adapter_id, adapter_name="turbo")
# pipe.load_lora_weights(adapter_id2, adapter_name="real")
# pipe.load_lora_weights(adapter_id3, weight_name="lora.safetensors", adapter_name="enhance")
# pipe.set_adapters(["turbo", "real", "enhance"], adapter_weights=[0.9, 0.66, 0.6])
# pipe.fuse_lora(adapter_names=["turbo", "real", "enhance"], lora_scale=1.0)
# pipe.unload_lora_weights()
# save to the Hub
# pipe.push_to_hub("FLUX.1M-8step_upscaler-cnet")
@spaces.GPU()
@torch.no_grad()
def caption(input_image: Image.Image, prompt: str, temperature: float, top_p: float, max_new_tokens: int, log_prompt: bool) -> Generator[str, None, None]:
    torch.cuda.empty_cache()
    if input_image is None:
        yield "No image provided. Please upload an image."
        return
    if log_prompt:
        print(f"PromptLog: {repr(prompt)}")
    convo = [
        {
            "role": "system",
            "content": "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions.",
        },
        {
            "role": "user",
            "content": prompt.strip(),
        },
    ]
    convo_string = cap_processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
    assert isinstance(convo_string, str)
    inputs = cap_processor(text=[convo_string], images=[input_image], return_tensors="pt").to('cuda')
    inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
    streamer = TextIteratorStreamer(cap_processor.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True if temperature > 0 else False,
        suppress_tokens=None,
        use_cache=True,
        temperature=temperature if temperature > 0 else None,
        top_k=None,
        top_p=top_p if temperature > 0 else None,
        streamer=streamer,
    )
    _ = cap_model.generate(**generate_kwargs)
    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)

        
@spaces.GPU()
@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)

    # Load control image
    control_image = load_image(control_image)
    w, h = control_image.size
    w = w - w % 32
    h = h - h % 32
    control_image = control_image.resize((int(w * scale), int(h * scale)), resample=2)  # Resample.BILINEAR
    print("Size to: " + str(control_image.size[0]) + ", " + str(control_image.size[1]))
    with torch.inference_mode():
        image = pipe(
            generator=generator,
            prompt=prompt,
            control_image=control_image,
            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]
        
    return image

# Create Gradio interface with rows and columns
with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
    gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
    with gr.Row():
        control_image = gr.Image(type="pil", label="Control Image", show_label=False)
        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, placeholder="Enter your prompt here...", label="Prompt")
            output_caption = gr.Textbox(label="Caption")
            scale = gr.Slider(1, 3, value=1, label="Scale", step=0.25)
            generate_button = gr.Button("Generate Image", variant="primary")
            caption_button = gr.Button("Generate Caption", variant="secondary")
        with gr.Column(scale=1):
            seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
            steps = gr.Slider(2, 16, value=8, label="Steps")
            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")
            
    
    with gr.Row():
        with gr.Accordion("Generation settings", open=False):
            prompt_box = 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.", visible=False)
            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
            )
            top_p_slider = gr.Slider(
                minimum=0.0, maximum=1.0, value=0.9, step=0.01,
                label="Top-p",
                visible=False
            )
            max_tokens_slider = gr.Slider(
                minimum=1, maximum=2048, value=512, 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
            )
        
        log_prompt = gr.Checkbox(value=True, label="Help improve JoyCaption by logging your text query", visible=False)
        
        gr.Markdown("**Tips:** 8 steps is all you need!")
    
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end],
        outputs=[generated_image]
    )
    caption_button.click(
		fn=caption,
		inputs=[control_image, prompt_box, temperature_slider, top_p_slider, max_tokens_slider, log_prompt],
		outputs=output_caption,
	)
# Launch the app
iface.launch()