diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index ff7c298861..c08a7ed634 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -324,6 +324,8 @@ dam = DescribeAnythingModel( )`, ]; +const diffusers_install = "pip install -U diffusers transformers"; + const diffusersDefaultPrompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"; const diffusersImg2ImgDefaultPrompt = "Turn this cat into a dog"; @@ -485,34 +487,37 @@ image = pipe( ]; export const diffusers = (model: ModelData): string[] => { + let codeSnippets: string[]; if ( model.tags.includes("StableDiffusionInpaintPipeline") || model.tags.includes("StableDiffusionXLInpaintPipeline") ) { - return diffusers_inpainting(model); + codeSnippets = diffusers_inpainting(model); } else if (model.tags.includes("controlnet")) { - return diffusers_controlnet(model); + codeSnippets = diffusers_controlnet(model); } else if (model.tags.includes("lora")) { if (model.pipeline_tag === "image-to-image") { - return diffusers_lora_image_to_image(model); + codeSnippets = diffusers_lora_image_to_image(model); } else if (model.pipeline_tag === "image-to-video") { - return diffusers_lora_image_to_video(model); + codeSnippets = diffusers_lora_image_to_video(model); } else if (model.pipeline_tag === "text-to-video") { - return diffusers_lora_text_to_video(model); + codeSnippets = diffusers_lora_text_to_video(model); } else { - return diffusers_lora(model); + codeSnippets = diffusers_lora(model); } } else if (model.tags.includes("textual_inversion")) { - return diffusers_textual_inversion(model); + codeSnippets = diffusers_textual_inversion(model); } else if (model.tags.includes("FluxFillPipeline")) { - return diffusers_flux_fill(model); + codeSnippets = diffusers_flux_fill(model); } else if (model.pipeline_tag === "image-to-video") { - return diffusers_image_to_video(model); + codeSnippets = diffusers_image_to_video(model); } else if (model.pipeline_tag === "image-to-image") { - return diffusers_image_to_image(model); + codeSnippets = diffusers_image_to_image(model); } else { - return diffusers_default(model); + codeSnippets = diffusers_default(model); } + + return [diffusers_install, ...codeSnippets]; }; export const diffusionkit = (model: ModelData): string[] => {