Update app.py

app.py

@@ -1,311 +1,139 @@
 import os
-import random
-import uuid
+from collections.abc import Iterator
+from threading import Thread
 import gradio as gr
-import numpy as np
-from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from typing import Tuple
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-css = '''
-.gradio-container{max-width: 575px !important}
-h1{text-align:center}
-footer {
-    visibility: hidden
-}
-'''
-
-DESCRIPTIONXX = """## TEXT 2 IMAGE🥠"""
-
-examples = [
+DESCRIPTION = """
+# QwQ Distill
+"""
 
-    "A tiny astronaut hatching from an egg on the moon, 4k, planet theme, --style raw5 --v 6.0",
-    "An anime-style illustration of a delicious, golden-brown wiener schnitzel on a plate, served with fresh lemon slices, parsley --style raw5",
-    "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5, 4K, Photo-Realistic",
-    "A cat holding a sign that says hello world --ar 85:128 --v 6.0 --style raw"
-]
+css= '''
+h1 {
+  text-align: center;
+  display: block;
+}
 
-MODEL_OPTIONS = {
-    
-    "LIGHTNING V5.0": "SG161222/RealVisXL_V5.0_Lightning",
-    "LIGHTNING V4.0": "SG161222/RealVisXL_V4.0_Lightning",
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
 }
+'''
 
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
-BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-style_list = [
-    {
-        "name": "3840 x 2160",
-        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-    },
-    {
-        "name": "2560 x 1440",
-        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-    },
-    {
-        "name": "HD+",
-        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-    },
-    {
-        "name": "Style Zero",
-        "prompt": "{prompt}",
-        "negative_prompt": "",
-    },
-]
-
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
-DEFAULT_STYLE_NAME = "3840 x 2160"
-STYLE_NAMES = list(styles.keys())
+model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.config.sliding_window = 4096
+model.eval()
 
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    if style_name in styles:
-        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    else:
-        p, n = styles[DEFAULT_STYLE_NAME]
 
-    if not negative:
-        negative = ""
-    return p.replace("{prompt}", positive), n + negative
-
-def load_and_prepare_model(model_id):
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        use_safetensors=True,
-        add_watermarker=False,
-    ).to(device)
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    
-    if USE_TORCH_COMPILE:
-        pipe.compile()
-    
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-    
-    return pipe
-
-# Preload and compile both models
-models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
-
-MAX_SEED = np.iinfo(np.int32).max
-
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-@spaces.GPU(duration=60, enable_queue=True)
+@spaces.GPU(duration=120)
 def generate(
-    model_choice: str,
-    prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    style_selection: str = DEFAULT_STYLE_NAME,
-    seed: int = 1,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    num_inference_steps: int = 25,
-    randomize_seed: bool = False,
-    use_resolution_binning: bool = True, 
-    num_images: int = 1,  
-    progress=gr.Progress(track_tqdm=True),
-):
-    global models
-    pipe = models[model_choice]
-    
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
-
-    options = {
-        "prompt": [prompt] * num_images,
-        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
-        "width": width,
-        "height": height,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "generator": generator,
-        "output_type": "pil",
-    }
-
-    if use_resolution_binning:
-        options["use_resolution_binning"] = True
-
-    images = []
-    for i in range(0, num_images, BATCH_SIZE):
-        batch_options = options.copy()
-        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
-        if "negative_prompt" in batch_options:
-            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
-        images.extend(pipe(**batch_options).images)
-
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
-
+    message: str,
+    chat_history: list[dict],
+    max_new_tokens: int = 1024,
+    temperature: float = 0.6,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
+) -> Iterator[str]:
+    conversation = chat_history.copy()
+    conversation.append({"role": "user", "content": message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    input_ids = input_ids.to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        {"input_ids": input_ids},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        repetition_penalty=repetition_penalty,
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown(DESCRIPTIONXX)
-    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.Gallery(label="Result", columns=1, show_label=False) 
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
 
-    with gr.Row():
-        model_choice = gr.Dropdown(
-            label="Model Selection⬇️",
-            choices=list(MODEL_OPTIONS.keys()),
-            value="LIGHTNING V5.0"
-        )
 
-    with gr.Accordion("Advanced options", open=False, visible=False):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=STYLE_NAMES,
-            value=DEFAULT_STYLE_NAME,
-            label="Quality Style",
-        )
-        num_images = gr.Slider(
-            label="Number of Images",
+demo = gr.ChatInterface(
+    fn=generate,
+    additional_inputs=[
+        gr.Slider(
+            label="Max new tokens",
             minimum=1,
-            maximum=5,
+            maximum=MAX_MAX_NEW_TOKENS,
             step=1,
-            value=1,
-        )
-        with gr.Row():
-            with gr.Column(scale=1):
-                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-                negative_prompt = gr.Text(
-                    label="Negative prompt",
-                    max_lines=5,
-                    lines=4,
-                    placeholder="Enter a negative prompt",
-                    value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-                    visible=True,
-                )
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Temperature",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
             step=1,
-            value=0,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row():
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=8,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=8,
-                value=1024,
-            )
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance Scale",
-                minimum=0.1,
-                maximum=6,
-                step=0.1,
-                value=3.0,
-            )
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps",
-                minimum=1,
-                maximum=60,
-                step=1,
-                value=28,
-            )
-
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        cache_examples=False
-    )
-
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        api_name=False,
-    )
-    
-    gr.on(
-        triggers=[
-            prompt.submit,
-            negative_prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            model_choice,
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            style_selection,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            randomize_seed, 
-            num_images, 
-        ],
-        outputs=[result, seed],
-    )
-
-    gr.Markdown(
-    """
-    <div style="text-align: justify;">
-    🥠Models used in the playground: <a href="https://huggingface.co/SG161222/RealVisXL_V5.0_Lightning">[LIGHTNING V5.0]</a>, <a href="https://huggingface.co/SG161222/RealVisXL_V4.0_Lightning">[LIGHTNING V4.0]</a>
-    for image generation. Stable Diffusion XL piped (SDXL) model HF. This is the demo space for generating images using the Stable Diffusion XL models, with multiple different variants available.
-    </div>
-    """
-    )
+            value=50,
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
+        ),
+    ],
+    stop_btn=None,
+    examples=[
+        ["Write a Python function to reverses a string if it's length is a multiple of 4. def reverse_string(str1): if len(str1) % 4 == 0: return ''.join(reversed(str1)) return str1 print(reverse_string('abcd')) print(reverse_string('python')) "],
+        ["Rectangle $ABCD$ is the base of pyramid $PABCD$. If $AB = 10$, $BC = 5$, $\overline{PA}\perp \text{plane } ABCD$, and $PA = 8$, then what is the volume of $PABCD$?"],
+        ["Difference between List comprehension and Lambda in Python lst  =  [x ** 2  for x in range (1, 11)   if  x % 2 == 1] print(lst)"],
+        ["What happens when the sun goes down?"],
+    ],
+    cache_examples=False,
+    type="messages",
+    description=DESCRIPTION,
+    css=css,
+    fill_height=True,
+)
 
-    gr.Markdown(
-    """
-    <div style="text-align: justify;">
-    🥠This is the demo space for generating images using Stable Diffusion XL with quality styles, different models, and types. Try the sample prompts to generate higher quality images. Try the sample prompts for generating higher quality images. 
-    <a href='https://huggingface.co/spaces/prithivMLmods/Top-Prompt-Collection' target='_blank'>Try prompts</a>.
-    </div>
-    """)
 
-    gr.Markdown(
-    """
-    <div style="text-align: justify;">
-    ⚠️ Users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.
-    </div>
-    """) 
-    
 if __name__ == "__main__":
-    demo.queue(max_size=50).launch(show_api=True)
+    demo.queue(max_size=20).launch()