Update app.py

app.py

@@ -4,7 +4,8 @@ import random
 import torch
 import spaces
 from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler
-from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+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"
@@ -36,19 +37,163 @@ examples = [
     "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.",
 ]
 
-css="""
+
+# 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"""# FLUX.1 [dev]
-12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
-[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
+        gr.Markdown(f"""TESTTESTTESTTESTTESTTEST]
         """)
         
         with gr.Row():
@@ -136,4 +281,7 @@ with gr.Blocks(css=css) as demo:
         outputs = [result, seed]
     )
 
-demo.launch()
+demo.launch()
+
+
+