running on t4 tests

app.py

@@ -52,30 +52,25 @@ def replace_concept_tokens(text: str):
     return text
 
 
-# def inference(
-#     prompt: str, num_inference_steps: int = 50, guidance_scale: int = 3.0
-# ):
-#     prompt = replace_concept_tokens(prompt)
-#     for _ in range(3):
-#         img_list = pipeline(
-#             prompt=prompt,
-#             num_inference_steps=num_inference_steps,
-#             guidance_scale=guidance_scale,
-#         )
-#         if not img_list["nsfw_content_detected"]:
-#             break
-#     return img_list["sample"]
+def inference(
+    prompt: str, guidance_scale: int, num_inference_steps: int, seed: int
+):
+    prompt = replace_concept_tokens(prompt)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    out = pipeline(
+        prompt=[prompt] * 2,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+    )
+    img_list = [item['sample'] for item in out]
+    return img_list
 
 DEFAULT_PROMPT = (
     "A watercolor painting on textured paper of a <det-logo> using soft strokes,"
     " pastel colors, incredible composition, masterpiece"
 )
 
-
-def white_imgs(prompt: str, guidance_scale: float, num_inference_steps: int, seed: int):
-    return [torch.ones(512, 512, 3).numpy() for _ in range(2)]
-
-
 with gr.Blocks() as demo:
     prompt = gr.Textbox(
         label="Prompt including the token '<det-logo>'",
@@ -110,11 +105,11 @@ with gr.Blocks() as demo:
     generate_btn = gr.Button(label="Generate")
     gallery = gr.Gallery(
         label="Generated Images",
-        value=[torch.zeros(512, 512, 3).numpy() for _ in range(2)],
+        value=[torch.ones(512, 512, 3).numpy() for _ in range(2)],
     ).style(height="auto")
 
     generate_btn.click(
-        white_imgs,
+        inference,
         inputs=[prompt, guidance_scale, num_inference_steps, seed],
         outputs=gallery,
     )