chore: load dart model

app.py

@@ -1,51 +1,115 @@
+import os
 import gradio as gr
 import numpy as np
 import random
-#import spaces #[uncomment to use ZeroGPU]
+import spaces
 from diffusers import DiffusionPipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
+try:
+    from dotenv import load_dotenv
+
+    load_dotenv()
+except:
+    print("failed to import dotenv (this is not a problem on the production)")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+HF_TOKEN = os.environ.get("HF_TOKEN")
+assert HF_TOKEN is not None
+
+IMAGE_MODEL_REPO_ID = os.environ.get(
+    "IMAGE_MODEL_REPO_ID", "OnomaAIResearch/Illustrious-xl-early-release-v0"
+)
+DART_V3_REPO_ID = os.environ.get("DART_V3_REPO_ID", "p1atdev/dart-v3-llama-8L-241003")
+
+torch_dtype = torch.bfloat16
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+dart = AutoModelForCausalLM.from_pretrained(
+    DART_V3_REPO_ID,
+    torch_dtype=torch_dtype,
+    token=HF_TOKEN,
+)
+tokenizer = AutoTokenizer.from_pretrained(DART_V3_REPO_ID)
+
+pipe = DiffusionPipeline.from_pretrained(IMAGE_MODEL_REPO_ID, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
+pipe = torch.compile(pipe)
+
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-#@spaces.GPU #[uncomment to use ZeroGPU]
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
+TEMPLATE = (
+    "<|bos|>"
+    #
+    "<|rating:general|>"
+    "{aspect_ratio}"
+    "<|length:medium|>"
+    #
+    "<copyright>original</copyright>"
+    #
+    "<character></character>"
+    #
+    "<general>"
+)
+
+
+@torch.inference_mode
+def generate_prompt(aspect_ratio: str):
+    input_ids = tokenizer.encode_plus(
+        TEMPLATE.format(aspect_ratio=aspect_ratio)
+    ).input_ids
 
+    output_ids = dart.generate(
+        input_ids,
+        max_new_tokens=256,
+        temperature=1.0,
+        top_p=1.0,
+        top_k=100,
+        num_beams=1,
+    )[0]
+
+    generated = output_ids[len(input_ids) :]
+    decoded = ", ".join(tokenizer.batch_decode(generated))
+
+    return decoded
+
+
+@spaces.GPU  # [uncomment to use ZeroGPU]
+def infer(
+    negative_prompt: str,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+    progress=gr.Progress(track_tqdm=True),
+):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-        
+
     generator = torch.Generator().manual_seed(seed)
-    
+
+    prompt = generate_prompt("<|aspect_ratio:square|>")
+    print(prompt)
+
     image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image, seed
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+    ).images[0]
+
+    return image, prompt, seed
+
+
+css = """
 #col-container {
     margin: 0 auto;
     max-width: 640px;
@@ -53,35 +117,28 @@ css="""
 """
 
 with gr.Blocks(css=css) as demo:
-    
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # Text-to-Image Gradio Template
+        # Random IllustriousXL
         """)
-        
+
         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)
-        
+            run_button = gr.Button("Generate random", scale=0)
+
         result = gr.Image(label="Result", show_label=False)
 
+        with gr.Accordion("Generation details", open=False):
+            prompt_txt = gr.Textbox("Generated prompt", interactive=False)
+
         with gr.Accordion("Advanced Settings", open=False):
-            
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
                 visible=False,
+                value=" worst quality, comic, multiple views, bad quality, low quality, lowres, displeasing, very displeasing, bad anatomy, bad hands, scan artifacts, monochrome, greyscale, signature, twitter username, jpeg artifacts, 2koma, 4koma, guro, extra digits, fewer digits",
             )
-            
+
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -89,54 +146,56 @@ with gr.Blocks(css=css) as demo:
                 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, #Replace with defaults that work for your model
+                    value=1024,  # Replace with defaults that work for your model
                 )
-                
+
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024, #Replace with defaults that work for your model
+                    value=1024,  # Replace with defaults that work for your model
                 )
-            
+
             with gr.Row():
-                
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
-                    minimum=0.0,
+                    minimum=1.0,
                     maximum=10.0,
-                    step=0.1,
-                    value=0.0, #Replace with defaults that work for your model
+                    step=0.5,
+                    value=6.5,
                 )
-                
+
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2, #Replace with defaults that work for your model
+                    value=20,
                 )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
+
     gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result, seed]
+        triggers=[run_button.click],
+        fn=infer,
+        inputs=[
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[result, prompt_txt, seed],
     )
 
-demo.queue().launch()
+demo.queue().launch()