fix: no pipe torch.compile

app.py

@@ -31,12 +31,15 @@ dart = AutoModelForCausalLM.from_pretrained(
     DART_V3_REPO_ID,
     torch_dtype=torch_dtype,
     token=HF_TOKEN,
+    use_cache=True,
 )
+dart = dart.eval()
+dart = dart.requires_grad_(False)
+dart = torch.compile(dart)
 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
@@ -60,12 +63,15 @@ TEMPLATE = (
 @torch.inference_mode
 def generate_prompt(aspect_ratio: str):
     input_ids = tokenizer.encode_plus(
-        TEMPLATE.format(aspect_ratio=aspect_ratio)
+        TEMPLATE.format(aspect_ratio=aspect_ratio),
+        return_tensors="pt",
     ).input_ids
+    print("input_ids", input_ids)
 
     output_ids = dart.generate(
         input_ids,
         max_new_tokens=256,
+        do_sample=True,
         temperature=1.0,
         top_p=1.0,
         top_k=100,
@@ -73,13 +79,35 @@ def generate_prompt(aspect_ratio: str):
     )[0]
 
     generated = output_ids[len(input_ids) :]
-    decoded = ", ".join(tokenizer.batch_decode(generated))
+    decoded = ", ".join([token for token in tokenizer.batch_decode(generated, skip_special_tokens=True) if token.strip() != ""])
+    print("decoded", decoded)
 
     return decoded
 
+@spaces.GPU
+def generate_image(
+    prompt: str,
+    negative_prompt: str,
+    generator,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    num_inference_steps: int,
+    progress=gr.Progress(track_tqdm=True),
+):
+    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
 
-@spaces.GPU  # [uncomment to use ZeroGPU]
-def infer(
+def on_generate(
     negative_prompt: str,
     seed,
     randomize_seed,
@@ -87,25 +115,23 @@ def infer(
     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]
+    image = generate_image(
+        prompt,
+        negative_prompt,
+        generator,
+        width,
+        height,
+        guidance_scale,
+        num_inference_steps,
+    )
 
     return image, prompt, seed
 
@@ -186,7 +212,7 @@ with gr.Blocks(css=css) as demo:
 
     gr.on(
         triggers=[run_button.click],
-        fn=infer,
+        fn=on_generate,
         inputs=[
             negative_prompt,
             seed,