app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)

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

@spaces.GPU(duration=190)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt = prompt, 
        width = width,
        height = height,
        num_inference_steps = num_inference_steps, 
        generator = generator,
        guidance_scale=guidance_scale
    ).images[0] 
    return image, seed
 
examples = [
    "a cat holding a sign that says hello world",
    "A scene full of classic video game characters as stickers on a black water bottle",
    "A futuristic biocity that is located in the former site of Portsmouth, New Hampshire. It has a mix of old and new buildings, green spaces, and water features. It also has six large artificial floating islands off of its coastline,(zenithal angle), ((by Iwan Baan)), coastal city,blue sky and white clouds,the sun is shining brightly,ultra-wide angle,",
    "Depict a breathtaking scene of a meteor rain showering down from a starry night sky. The meteors should vary in size and brightness, streaking across the sky with vibrant tails of light, creating a dazzling display. Below, a serene landscape—perhaps a tranquil lake reflecting the celestial spectacle, or a rugged mountain range—should enhance the sense of wonder. The foreground can include silhouettes of trees or figures gazing up in awe at the cosmic event. The overall atmosphere should evoke feelings of magic and inspiration, capturing the beauty and mystery of the universe.",
]


# Image Gen css - Only saving for backup - ***disregard***
#css="""
#col-container {
#    margin: 0 auto;
#    max-width: 520px;
#}
#"""

# Llama + Flux CSS
CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
.chatbox .messages .message.user {
    background-color: #e1f5fe;
}
.chatbox .messages .message.bot {
    background-color: #eeeeee;
}
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""
#################################################################################################################################
##########################################


model_id = "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
  model_id,
  torch_dtype=torch.float16,
  device_map="sequential",
  offload_folder="offload",
  offload_state_dict=True 
)

TITLE = "Quick Description"

DESCRIPTION = """
Generate a longer description for your image from a simple basic prompt
"""



@spaces.GPU(duration=120)
def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
    print(f'Message: {message}')
    print(f'History: {history}')
    
    conversation = []
    for prompt, answer in history:
        conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
    conversation.append({"role": "user", "content": message})
    
    input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_ids, return_tensors="pt").to(0)
    
    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = dict(
        inputs, 
        streamer=streamer,
        top_k=top_k,
        top_p=top_p,
        repetition_penalty=penalty,
        max_new_tokens=max_new_tokens, 
        do_sample=True, 
        temperature=temperature,
        eos_token_id=[128001, 128009],
    )
    
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer

chatbot = gr.Chatbot(height=500)

with gr.Blocks(css=CSS) as demo:
    gr.HTML(TITLE)
    gr.HTML(DESCRIPTION)
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        theme="soft",
        retry_btn=None,
        undo_btn="Delete Previous",
        clear_btn="Clear",
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.8,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=128,
                maximum=4096,
                step=1,
                value=1024,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=0.8,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=1,
                maximum=20,
                step=1,
                value=20,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=1.2,
                label="Repetition penalty",
                render=False,
            ),
        ],
        examples=[
            ["Explain Deep Learning as a pirate."],
            ["Give me five ideas for a child's summer science project."],
            ["Provide advice for writing a script for a puzzle game."],
            ["Create a tutorial for building a breakout game using markdown."]
        ],
        cache_examples=False,
    )
#################################################################################################################################
with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""TESTTESTTESTTESTTESTTEST]
        """)
        
        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)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            
            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=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
        gr.Examples(
            examples = examples,
            fn = infer,
            inputs = [prompt],
            outputs = [result, seed],
            cache_examples="lazy"
        )

        def greet(name):
            return "Hello " + name + "! Imagine an image with Flux"

        name = gr.Textbox(label="Name")
        output = gr.Textbox(label="Output Box")
        greet_btn = gr.Button("Greet")
        greet_btn.click(fn=greet, inputs=name, outputs=output, api_name="greet")

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs = [result, seed]
    )

demo.launch()